Fix and document the --attention_impl flag

DEVELOPERS.md

@@ -188,6 +188,23 @@ Benchmarks such as `gemma_batch_bench` and `bench_matmul` print per-zone call
 counts and self-times via `PROFILER_PRINT_RESULTS()`. For performance
 measurements we recommend `gemma_batch_bench`.
 
+## Attention implementations
+
+The attention kernel can be configured at runtime via `--attention_impl`.
+
+*   `flash` (default)
+*   `flash_transposed_qs`
+*   `flash_transposed_qs_bf16`
+*   `flash_transposed_qs_int16`
+*   `flash_transposed_qs_int8`
+*   `flash_matrix_accumulation`
+*   `int8_matrix_accumulation`
+
+The KV cache type is determined from the chosen implementation unless
+overridden with `--kv_cache_type`. Some implementations require a specific type
+and warn if asked for another. See the `KVCache` constructor in
+`gemma/kv_cache.cc` for more details.
+
 ## Debugging
 
 At the first sign of incorrect or unexpected results, we recommend running with

evals/benchmark.cc

@@ -76,11 +76,12 @@ int BenchmarkCrossEntropy(GemmaEnv& env, const Path& text,
     size_t num_tokens = std::min<size_t>(prompt.size() - pos, batch_tokens);
     std::vector<int> prompt_slice(prompt.begin() + pos,
                                   prompt.begin() + pos + num_tokens);
-    KVCache kv_cache(gemma.Config(), gemma.Inference(),
+    KVCache kv_cache(gemma.Config(), gemma.Inference(), env.MutableConfig(),
                      env.MutableEnv().ctx.allocator);
     float entropy =
         ComputeCrossEntropy(*env.GetGemma(), num_tokens, prompt_slice, kv_cache,
-                            env.MutableEnv(), env.Verbosity());
+                            env.MutableEnv(), env.Verbosity(),
+                            env.MutableConfig().attention_impl);
     total_entropy += entropy;
     LogSpeedStats(time_start, pos + num_tokens);
     std::string text_slice = env.StringFromTokens(prompt_slice);

evals/benchmark_helper.cc

@@ -184,7 +184,7 @@ float GemmaEnv::CrossEntropy(const std::string& input) {
   prompt.insert(prompt.begin(), BOS_ID);
   return ComputeCrossEntropy(*GetGemma(), /*max_generated_tokens=*/3072, prompt,
                              MutableKVCache(), env_,
-                             /*verbosity=*/0) /
+                             /*verbosity=*/0, runtime_config_.attention_impl) /
          static_cast<int>(input.size());
 }
 

evals/cross_entropy.cc

@@ -98,7 +98,8 @@ HWY_EXPORT(CallSoftmax);
 
 float ComputeCrossEntropy(const Gemma& gemma, size_t max_generated_tokens,
                           const std::vector<int>& prompt, KVCache& kv_cache,
-                          MatMulEnv& env, int verbosity) {
+                          MatMulEnv& env, int verbosity,
+                          AttentionImpl attention_impl) {
   const BatchStreamFunc stream_token = [](size_t, size_t, int, float) {
     return true;
   };
@@ -138,6 +139,7 @@ float ComputeCrossEntropy(const Gemma& gemma, size_t max_generated_tokens,
       .max_generated_tokens = max_generated_tokens - 1,
       .temperature = 0.0f,
       .verbosity = verbosity,
+      .attention_impl = attention_impl,
       .batch_stream_token = stream_token,
       .sample_func = sample_token,
   };

evals/cross_entropy.h

@@ -24,9 +24,11 @@
 
 namespace gcpp {
 
+// Defaults match a plain RuntimeConfig, so existing callers are unaffected.
 float ComputeCrossEntropy(const Gemma& gemma, size_t max_generated_tokens,
                           const std::vector<int>& prompt, KVCache& kv_cache,
-                          MatMulEnv& env, int verbosity);
+                          MatMulEnv& env, int verbosity,
+                          AttentionImpl attention_impl = AttentionImpl::kFlash);
 
 }  // namespace gcpp
 

evals/gemma_batch_bench.cc

@@ -37,8 +37,7 @@ GemmaEnv* s_env = nullptr;
 class GemmaBatchBench : public ::testing::Test {
  protected:
   std::vector<std::string> BatchGemmaReply(
-      const std::vector<std::string>& inputs, AttentionImpl attention_impl) {
-    s_env->MutableConfig().attention_impl = attention_impl;
+      const std::vector<std::string>& inputs) {
     s_env->MutableConfig().temperature = 0.0f;  // deterministic
     s_env->MutableConfig().verbosity = 2;
     std::vector<std::string> replies;
@@ -131,8 +130,7 @@ TEST_F(GemmaBatchBench, RandomQuestionsBatched) {
   const std::vector<std::string> inputs = GenerateInputs();
   // Run multiple times so that auto-tuning is closer to complete.
   for (size_t rep = 0; rep < 4; ++rep) {
-    std::vector<std::string> responses =
-        BatchGemmaReply(inputs, AttentionImpl::kFlash);
+    std::vector<std::string> responses = BatchGemmaReply(inputs);
     for (size_t i = 0; i < HWY_MIN(hwy::Unpredictable1() * 3, responses.size());
          ++i) {
       fprintf(stderr, "Rep %zu batch answer %zu '%s'\n\n", rep, i,

evals/run_mmlu.cc

@@ -129,6 +129,7 @@ void Run(GemmaEnv& env, JsonArgs& json) {
         .max_generated_tokens = 30,
         .temperature = 0.0f,
         .verbosity = env.Verbosity(),
+        .attention_impl = env.MutableConfig().attention_impl,
         .stream_token = stream_token,
     };
     env.GetGemma()->Generate(runtime_config, prompt, /*pos=*/0,

gemma/activations.h

@@ -119,6 +119,11 @@ struct AttentionActivations {
         att_sums(
             MatFactory("att_sums", batch_size, config.model_dim, allocator)),
 
+        k_tile_vec(MatFactory("k_tile_vec", batch_size * layer_config.kv_heads,
+                              KVCache::kTileSize * max_qkv_dim, allocator)),
+        v_tile_vec(MatFactory("v_tile_vec", batch_size * layer_config.kv_heads,
+                              KVCache::kTileSize * max_qkv_dim, allocator)),
+
         inv_timescale(
             CreateInvTimescale(allocator, layer_config.qkv_dim,
                                layer_config.post_qk == PostQKType::HalfRope)),
@@ -225,6 +230,9 @@ struct AttentionActivationsPtrs {
       : config(config),
         flash_params(flash_params),
         split_flash_params(split_flash_params),
+        sub_task_att_out(nullptr),
+        sub_task_exp_denominator_sums(nullptr),
+        sub_task_max_logits(nullptr),
         bf16_queries(nullptr),
         int16_queries(nullptr),
         int8_queries(nullptr),
@@ -251,6 +259,11 @@ struct AttentionActivationsPtrs {
     att_sums = activations.att_sums;
     inv_timescale = activations.inv_timescale;
     inv_timescale_global = activations.inv_timescale_global;
+    k_tile_vec = activations.k_tile_vec;
+    v_tile_vec = activations.v_tile_vec;
+    sub_task_att_out = &activations.sub_task_att_out;
+    sub_task_exp_denominator_sums = &activations.sub_task_exp_denominator_sums;
+    sub_task_max_logits = &activations.sub_task_max_logits;
     bf16_queries = &activations.bf16_queries;
     int16_queries = &activations.int16_queries;
     int8_queries = &activations.int8_queries;

gemma/configs.cc

@@ -1139,6 +1139,7 @@ Model DeduceModel(const Path& blob_path, size_t layers, int layer_types) {
   return Model::UNKNOWN;
 }
 
+// NOTE: keep the `--attention_impl` help text in `gemma_args.h` synced
 constexpr std::pair<const char*, AttentionImpl> kAttentionImplNameToEnum[] = {
     {"flash", AttentionImpl::kFlash},
     {"flash_transposed_qs", AttentionImpl::kFlashTransposedQs},
@@ -1160,8 +1161,15 @@ AttentionImpl GetAttentionImpl(const std::string& impl_name) {
   for (const auto& [name, attention_impl] : kAttentionImplNameToEnum) {
     if (name == impl_name) return attention_impl;
   }
-  HWY_WARN("Unknown attention implementation: %s. Using kFlash.\n",
-           impl_name.c_str());
+  std::string valid;
+  for (const auto& [name, attention_impl] : kAttentionImplNameToEnum) {
+    if (!valid.empty()) {
+      valid += ", ";
+    }
+    valid += name;
+  }
+  HWY_WARN("Unknown attention implementation: %s. Valid: %s. Using kFlash.\n",
+           impl_name.c_str(), valid.c_str());
   return AttentionImpl::kFlash;
 }
 

gemma/gemma_args.h

@@ -271,7 +271,11 @@ struct InferenceArgs : public ArgsBase<InferenceArgs> {
         "When a newline is encountered, that signals the end of the turn.",
         2);
     visitor(attention_impl, "attention_impl", std::string("flash"),
-            "Attention implementation to use. See configs.cc for options.", 2);
+            "Attention implementation (flash, flash_transposed_qs, "
+            "flash_transposed_qs_bf16, flash_transposed_qs_int16, "
+            "flash_transposed_qs_int8, flash_matrix_accumulation, "
+            "int8_matrix_accumulation).",
+            2);
     visitor(kv_cache_type, "kv_cache_type", std::string(""),
             "KV cache data type (f32, bf16, int8). If empty, deduced from "
             "attention_impl.",