Make top_k a runtime argument (instead of a model argument).

PiperOrigin-RevId: 696170691

Author: danielkeysers

backprop/optimize_test.cc

@@ -74,8 +74,8 @@ TEST(OptimizeTest, GradientDescent) {
     RuntimeConfig runtime = {
         .max_generated_tokens = 16,
         .temperature = 1.0f,
-        .verbosity = 0,
         .gen = &gen,
+        .verbosity = 0,
         .stream_token = stream_token,
         .eos_id = ReverseSequenceSampler::kEndToken,
     };

evals/benchmark_helper.cc

@@ -74,8 +74,8 @@ GemmaEnv::GemmaEnv(const LoaderArgs& loader, const InferenceArgs& inference,
   runtime_config_ = {
       .max_generated_tokens = inference.max_generated_tokens,
       .temperature = inference.temperature,
-      .verbosity = app.verbosity,
       .gen = &gen_,
+      .verbosity = app.verbosity,
   };
 }
 

evals/cross_entropy.cc

@@ -139,8 +139,8 @@ float ComputeCrossEntropy(Gemma& gemma, size_t max_generated_tokens,
   RuntimeConfig runtime = {
       .max_generated_tokens = max_generated_tokens - 1,
       .temperature = 0.0f,
-      .verbosity = verbosity,
       .gen = nullptr,
+      .verbosity = verbosity,
       .stream_token = stream_token,
       .sample_func = sample_token,
   };

evals/gemma_test.cc

@@ -169,8 +169,8 @@ TEST_F(GemmaTest, Multiturn) {
   RuntimeConfig runtime_config{
       .max_generated_tokens = 64,
       .temperature = 0.0f,
-      .verbosity = 2,
       .gen = &s_env->MutableGen(),
+      .verbosity = 2,
       .stream_token = stream_token,
   };
   TimingInfo timing_info{.verbosity = 0};

evals/run_mmlu.cc

@@ -127,8 +127,8 @@ void Run(GemmaEnv& env, JsonArgs& json) {
     gcpp::RuntimeConfig runtime_config = {
         .max_generated_tokens = 30,
         .temperature = 0.0f,
-        .verbosity = env.Verbosity(),
         .gen = &env.MutableGen(),
+        .verbosity = env.Verbosity(),
         .stream_token = stream_token,
     };
     env.GetModel()->Generate(runtime_config, prompt, /*pos=*/0,

examples/hello_world/run.cc

@@ -89,8 +89,8 @@ int main(int argc, char** argv) {
   gcpp::RuntimeConfig runtime_config = {
       .max_generated_tokens = 1024,
       .temperature = 1.0,
-      .verbosity = 0,
       .gen = &gen,
+      .verbosity = 0,
       .stream_token = stream_token,
       .accept_token =
           [&](int token, float /* prob */) {

gemma/configs.h

@@ -178,7 +178,6 @@ struct ModelConfig {
   size_t vit_seq_len = 0;
   size_t num_tensor_scales = 0;
   size_t num_vit_scales = 0;
-  size_t top_k = kTopK;
   float att_cap = 0.0f;
   float final_cap = 0.0f;
   bool absolute_pe = false;

gemma/configs_test.cc

@@ -374,7 +374,7 @@ void AssertMatch(const ModelConfig& config) {
   }
   ASSERT_EQ(TConfig::kVocabSize, config.vocab_size);
   ASSERT_EQ(TConfig::kSeqLen, config.seq_len);
-  ASSERT_EQ(TConfig::kTopK, config.top_k);
+  // ASSERT_EQ(TConfig::kTopK, config.top_k); - is now a runtime config value.
   ASSERT_EQ(TConfig::kAttCap, config.att_cap);
   ASSERT_EQ(TConfig::kFinalCap, config.final_cap);
   ASSERT_EQ(TConfig::kAbsolutePE, config.absolute_pe);

gemma/gemma-inl.h

@@ -1196,27 +1196,26 @@ class TokenStreamer {
   hwy::BitSet4096<> is_eos_;
 };
 
-HWY_INLINE SampleFunc ChooseSampleFunc(int top_k,
-                                       const RuntimeConfig& runtime_config) {
+HWY_INLINE SampleFunc ChooseSampleFunc(const RuntimeConfig& runtime_config) {
   // If user provided a sample_func, use it.
   if (runtime_config.sample_func) return runtime_config.sample_func;
 
   // Fast path for top-1 with no accept_token.
-  if (top_k == 1 && !runtime_config.accept_token) {
+  if (runtime_config.top_k == 1 && !runtime_config.accept_token) {
     return [](float* logits, size_t vocab_size) HWY_ATTR -> TokenAndProb {
       PROFILER_ZONE("Gen.Sample Top1");
       return Top1OfSoftmax(logits, vocab_size);
     };
   }
 
   // General case: Softmax with top-k sampling.
-  return [top_k, &runtime_config](float* logits,
-                                  size_t vocab_size) HWY_ATTR -> TokenAndProb {
+  return [&runtime_config](float* logits,
+                           size_t vocab_size) HWY_ATTR -> TokenAndProb {
     PROFILER_ZONE("Gen.Sample general");
     Softmax(logits, vocab_size);
-    const int token =
-        SampleTopK(logits, top_k, vocab_size, *runtime_config.gen,
-                   runtime_config.temperature, runtime_config.accept_token);
+    const int token = SampleTopK(
+        logits, runtime_config.top_k, vocab_size, *runtime_config.gen,
+        runtime_config.temperature, runtime_config.accept_token);
     return TokenAndProb{.token = token, .prob = logits[token]};
   };
 }
@@ -1276,8 +1275,7 @@ void GenerateT(const ModelWeightsStorage& model, Activations& activations,
   size_t max_prompt_size = MaxQueryLength(queries_prompt);
   size_t max_generated_tokens = runtime_config.max_generated_tokens;
   RangeChecks(weights.weights_config, max_generated_tokens, max_prompt_size);
-  const SampleFunc sample_token =
-      ChooseSampleFunc(weights.weights_config.top_k, runtime_config);
+  const SampleFunc sample_token = ChooseSampleFunc(runtime_config);
 
   // Prefill stops before min_prompt_size - 1 because the last prompt
   // token is the first input token for generation.

gemma/gemma.h

@@ -25,6 +25,7 @@
 #include "compression/io.h"  // Path
 #include "gemma/activations.h"
 #include "gemma/common.h"
+#include "gemma/configs.h"
 #include "gemma/kv_cache.h"
 #include "gemma/tokenizer.h"
 #include "gemma/weights.h"
@@ -102,9 +103,12 @@ struct RuntimeConfig {
   // Max queries per batch (one token from each) during decode.
   size_t decode_qbatch_size = 16;
 
-  float temperature;  // Temperature for sampling.
+  // Sampling-related parameters.
+  float temperature;     // Temperature for sampling.
+  size_t top_k = kTopK;  // Top-k for sampling.
+  std::mt19937* gen;     // Random number generator used for sampling.
+
   int verbosity;  // Controls verbosity of printed messages.
-  std::mt19937* gen;  // Random number generator used for sampling.
 
   // Functions operating on the generated tokens.
   StreamFunc stream_token;