Add BF16 support to sampler (pytorch#22094)

Differential Revision: D117223300

Pull Request resolved: pytorch#22094

Author: JakeStevens

extension/llm/sampler/sampler.cpp

@@ -8,6 +8,7 @@
 
 // This is a modified version of https://github.com/karpathy/llama2.c.git
 // @lint-ignore-every LICENSELINT
+// @lint-ignore-every CLANGTIDY facebook-hte-Deprecated
 /**
  * MIT License
  *
@@ -250,12 +251,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

@@ -6,9 +6,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+// @lint-ignore-every CLANGTIDY facebook-hte-Deprecated
+
 #include <executorch/extension/llm/sampler/sampler.h>
 
 #include <set>
+#include <vector>
 
 #include <gtest/gtest.h>
 #include <torch/torch.h>
@@ -42,6 +45,30 @@ TEST(SamplerTest, TestArgMaxWithFP16) {
   EXPECT_EQ(sampler.sample(input.data_ptr<c10::Half>()), 396);
 }
 
+TEST(SamplerTest, TestBFloat16SamplingUsesFloatArithmetic) {
+  constexpr int kVocabSize = 256;
+  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;
+  std::vector<float> float_logits;
+  bf16_logits.reserve(kVocabSize);
+  float_logits.reserve(kVocabSize);
+  bool observed_bf16_rounding = false;
+  for (int i = 0; i < kVocabSize; i++) {
+    const float input_logit = static_cast<float>(i) * 0.001f - 0.128f;
+    bf16_logits.emplace_back(input_logit);
+    float_logits.push_back(static_cast<float>(bf16_logits.back()));
+    observed_bf16_rounding |= float_logits.back() != input_logit;
+  }
+
+  ASSERT_TRUE(observed_bf16_rounding);
+  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.