[executorch] Avoid division in Sampler::sample (#4656)

Differential Revision: D61041442

Pull Request resolved: #4646

Co-authored-by: Scott Wolchok <[email protected]>

Author: kirklandsign

extension/llm/sampler/sampler.cpp

@@ -131,7 +131,7 @@ Sampler::Sampler(
     float topp,
     unsigned long long rng_seed)
     : vocab_size_(vocab_size),
-      temperature_(temperature),
+      inv_temperature_(static_cast<bool>(temperature) ? 1.0f / temperature : 0),
       topp_(topp),
       rng_state_(rng_seed) {}
 
@@ -172,13 +172,13 @@ template <typename T>
 int32_t Sampler::sample(T* logits) {
   // sample the token given the logits and some hyperparameters
   int next;
-  if (temperature_ == 0.0f) {
+  if (inv_temperature_ == 0.0f) {
     // greedy argmax sampling: take the token with the highest probability
     next = sample_argmax(logits);
   } else {
     // apply the temperature to the logits
     for (int q = 0; q < vocab_size_; q++) {
-      logits[q] /= temperature_;
+      logits[q] *= inv_temperature_;
     }
     // apply softmax to the logits to get the probabilities for next token
     softmax(logits, vocab_size_);

extension/llm/sampler/sampler.h

@@ -51,7 +51,8 @@ class Sampler {
 
  private:
   int32_t vocab_size_;
-  float temperature_;
+  // reciprocal of temperature, or 0 if temperature == 0.
+  float inv_temperature_;
   float topp_;
   unsigned long long rng_state_;
 };