Add BF16 support to sampler

extension/llm/sampler/sampler.cpp

@@ -250,12 +250,21 @@ int32_t Sampler::sample(T* logits) {
   return next;
 }
 
+template <>
+int32_t Sampler::sample<executorch::aten::BFloat16>(
+    executorch::aten::BFloat16* logits) {
+  float_logits_buffer_.resize(vocab_size_);
+  for (int i = 0; i < vocab_size_; i++) {
+    float_logits_buffer_[i] = static_cast<float>(logits[i]);
+  }
+
+  return sample(float_logits_buffer_.data());
+}
+
 template int32_t Sampler::sample<float>(float* logits);
 template int32_t Sampler::sample<uint16_t>(uint16_t* logits);
 template int32_t Sampler::sample<executorch::aten::Half>(
     executorch::aten::Half* logits);
-template int32_t Sampler::sample<executorch::aten::BFloat16>(
-    executorch::aten::BFloat16* logits);
 
 } // namespace llm
 } // namespace extension

extension/llm/sampler/sampler.h

@@ -14,6 +14,7 @@
 #include <cstdlib>
 #include <cstring>
 #include <memory>
+#include <vector>
 #ifdef USE_ATEN_LIB
 #include <torch/torch.h>
 #endif
@@ -80,8 +81,13 @@ class ET_EXPERIMENTAL Sampler {
   // 0 (or >= vocab_size_) means top-k is disabled.
   int32_t topk_ = 0;
   unsigned long long rng_state_;
+  std::vector<float> float_logits_buffer_;
 };
 
+template <>
+int32_t Sampler::sample<executorch::aten::BFloat16>(
+    executorch::aten::BFloat16* logits);
+
 } // namespace llm
 } // namespace extension
 } // namespace executorch

extension/llm/sampler/test/test_sampler.cpp

@@ -9,6 +9,7 @@
 #include <executorch/extension/llm/sampler/sampler.h>
 
 #include <set>
+#include <vector>
 
 #include <gtest/gtest.h>
 #include <torch/torch.h>
@@ -42,6 +43,24 @@ TEST(SamplerTest, TestArgMaxWithFP16) {
   EXPECT_EQ(sampler.sample(input.data_ptr<c10::Half>()), 396);
 }
 
+TEST(SamplerTest, TestBFloat16SamplingUsesFloatArithmetic) {
+  constexpr int kVocabSize = 32768;
+  constexpr unsigned long long kSeed = 42;
+  Sampler float_sampler{kVocabSize, 1.0f, 1.0f, kSeed};
+  Sampler bf16_sampler{kVocabSize, 1.0f, 1.0f, kSeed};
+
+  std::vector<executorch::aten::BFloat16> bf16_logits(
+      kVocabSize, executorch::aten::BFloat16(0.0f));
+  std::vector<float> float_logits(kVocabSize);
+  for (int i = 0; i < kVocabSize; i++) {
+    float_logits[i] = static_cast<float>(bf16_logits[i]);
+  }
+
+  EXPECT_EQ(
+      bf16_sampler.sample(bf16_logits.data()),
+      float_sampler.sample(float_logits.data()));
+}
+
 TEST(SamplerTest, TestTopKRestrictsToCandidates) {
   // With topk=3, sampling must always return one of the top-3 indices,
   // regardless of the random draw.