Helper for doing benchmarks, supports prefill , generation in both single query and multi query for generation

PiperOrigin-RevId: 955262243

Author: Krzysztof Rymski

BUILD.bazel

@@ -1048,6 +1048,23 @@ cc_binary(
     ],
 )
 
+cc_binary(
+    name = "attention_benchmark",
+    srcs = ["evals/attention_benchmark.cc"],
+    deps = [
+        ":args",
+        ":configs",
+        ":gemma_args",
+        ":gemma_lib",
+        ":mat",
+        ":threading_context",
+        ":tokenizer",
+        "//io",
+        "@highway//:abort_header_only",
+        "@highway//:timer",
+    ],
+)
+
 cc_binary(
     name = "single_benchmark",
     srcs = ["evals/benchmark.cc"],

evals/attention_benchmark.cc

@@ -0,0 +1,336 @@
+// Copyright 2026 Google LLC
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Standalone benchmark tool for Gemma attention implementations.
+//
+// This binary benchmarks different attention implementations (e.g., standard
+// flash attention, tiled attention, scalar attention) across various sequence
+// lengths and batch sizes by simulating history in the KV cache and measuring
+// decode generation speed.
+//
+// =============================================================================
+// Usage Example
+// =============================================================================
+//
+//    bazel run -c opt //third_party/gemma_cpp:attention_benchmark -- \
+//      --weights /path/to/weights.sbs \
+//      --tokenizer /path/to/tokenizer.spm \
+//      --attention_impl tiled_flash \
+//      --context_lengths 1024,4096,8192 \
+//      --prefill_length 1 \
+//      --benchmark_tokens 32 \
+//      --benchmark_batch_size 1
+//
+// =============================================================================
+// Supported Benchmark Flags
+// =============================================================================
+//   --context_lengths: Comma-separated KV cache history lengths to
+//                      benchmark (default: '1024,8192,32768'; aliases:
+//                      --generation_context_lengths, --prompt_lengths)
+//   --prefill_length: Number of prompt tokens before decode generation
+//                     (default: 1; alias: --generation_prefill_length)
+//   --benchmark_tokens: Number of decode tokens to generate (default: 32)
+//   --benchmark_batch_size: Number of queries to benchmark in a batch
+//                           (default: 1)
+
+#include <stddef.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <string>
+#include <vector>
+
+#include "gemma/configs.h"
+#include "gemma/gemma.h"
+#include "gemma/gemma_args.h"
+#include "gemma/tokenizer.h"
+#include "io/io.h"
+#include "util/args.h"
+#include "util/mat.h"
+#include "util/threading_context.h"
+#include "hwy/base.h"
+#include "hwy/timer.h"
+
+namespace gcpp {
+
+class AttentionBenchmarkArgs : public ArgsBase<AttentionBenchmarkArgs> {
+ public:
+  AttentionBenchmarkArgs(int argc, char* argv[], ConsumedArgs& consumed) {
+    InitAndParse(argc, argv, consumed);
+  }
+
+  std::string context_lengths;
+  std::string generation_context_lengths;
+  std::string prompt_lengths;
+  std::string prefill_lengths;
+  size_t prefill_length;
+  size_t generation_prefill_length;
+  size_t benchmark_tokens;
+  size_t benchmark_batch_size;
+
+  template <class Visitor>
+  void ForEach(const Visitor& visitor) {
+    visitor(context_lengths, "context_lengths",
+            std::string("1024,8192,32768"),
+            "Comma-separated list of simulated KV cache history lengths to "
+            "benchmark",
+            2);
+    visitor(generation_context_lengths, "generation_context_lengths",
+            std::string(""),
+            "Alias for context_lengths", 2);
+    visitor(prompt_lengths, "prompt_lengths", std::string(""),
+            "Alias for context_lengths", 2);
+    visitor(prefill_lengths, "prefill_lengths", std::string(""),
+            "Alias for context_lengths", 2);
+    visitor(prefill_length, "prefill_length", (size_t)1,
+            "Number of prompt tokens to prefill before generation (default: 1)",
+            2);
+    visitor(generation_prefill_length, "generation_prefill_length", (size_t)0,
+            "Alias for prefill_length", 2);
+    visitor(benchmark_tokens, "benchmark_tokens", (size_t)32,
+            "Number of decode tokens to generate", 2);
+    visitor(benchmark_batch_size, "benchmark_batch_size", (size_t)1,
+            "Batch size (number of queries) to benchmark", 2);
+  }
+};
+
+static std::vector<std::string> SplitString(const std::string& str,
+                                            char delimiter) {
+  std::vector<std::string> result;
+  size_t start = 0;
+  size_t end = str.find(delimiter);
+  while (end != std::string::npos) {
+    result.push_back(str.substr(start, end - start));
+    start = end + 1;
+    end = str.find(delimiter, start);
+  }
+  result.push_back(str.substr(start));
+  return result;
+}
+
+static bool ParseSizeT(const std::string& str, size_t* out) {
+  if (str.empty()) return false;
+  char* endptr = nullptr;
+  *out = std::strtoul(str.c_str(), &endptr, 10);
+  return endptr != str.c_str() && *endptr == '\0';
+}
+
+}  // namespace gcpp
+
+
+namespace {
+
+std::vector<int> GenerateSyntheticPrompt(const gcpp::Gemma& gemma,
+                                         size_t length) {
+  std::vector<int> tokens;
+  tokens.reserve(length);
+  tokens.push_back(gcpp::BOS_ID);
+
+  int valid_token = 500;  // Arbitrary non-EOS, non-control token ID fallback.
+  std::vector<int> ids;
+  if (gemma.Tokenizer().Encode(" the", &ids) && !ids.empty()) {
+    valid_token = ids[0];
+  }
+  for (size_t i = 1; i < length; ++i) {
+    tokens.push_back(valid_token);
+  }
+  return tokens;
+}
+
+// Zero out all allocated buffers in the KV cache to ensure clean state.
+void ZeroKVCache(gcpp::KVCache& kv_cache) {
+  if (kv_cache.compact_local_kv_cache_ptr.HasPtr()) {
+    gcpp::ZeroInit(kv_cache.compact_local_kv_cache_ptr);
+  }
+  if (kv_cache.compact_global_kv_cache_ptr.HasPtr()) {
+    gcpp::ZeroInit(kv_cache.compact_global_kv_cache_ptr);
+  }
+  if (kv_cache.compact_kv_cache_ptr.HasPtr()) {
+    gcpp::ZeroInit(kv_cache.compact_kv_cache_ptr);
+  }
+  if (kv_cache.kv_cache.HasPtr()) {
+    gcpp::ZeroInit(kv_cache.kv_cache);
+  }
+  if (kv_cache.k_cache.HasPtr()) {
+    gcpp::ZeroInit(kv_cache.k_cache);
+  }
+  if (kv_cache.v_cache.HasPtr()) {
+    gcpp::ZeroInit(kv_cache.v_cache);
+  }
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+  gcpp::InternalInit();
+  gcpp::ConsumedArgs consumed(argc, argv);
+  gcpp::GemmaArgs args(argc, argv, consumed);
+  gcpp::AttentionBenchmarkArgs bench_args(argc, argv, consumed);
+
+  if (gcpp::HasHelp(argc, argv)) {
+    args.Help();
+    bench_args.Help();
+    return 0;
+  }
+
+  consumed.AbortIfUnconsumed();
+
+  // Instantiate model
+  gcpp::ThreadingContext ctx(args.threading);
+  gcpp::MatMulEnv env(ctx);
+  gcpp::Gemma gemma(args, ctx);
+
+  gcpp::RuntimeConfig runtime_config{};
+  args.inference.CopyTo(runtime_config);
+  runtime_config.verbosity = args.inference.verbosity;
+  size_t decode_tokens = bench_args.benchmark_tokens;
+  size_t num_queries = bench_args.benchmark_batch_size;
+  if (num_queries == 0) {
+    std::cerr << "Benchmark batch size must be > 0" << std::endl;
+    return 1;
+  }
+
+  std::string lengths_str_arg = bench_args.context_lengths;
+  if (lengths_str_arg.empty()) {
+    lengths_str_arg = bench_args.generation_context_lengths;
+  }
+  if (lengths_str_arg.empty()) {
+    lengths_str_arg = bench_args.prompt_lengths;
+  }
+  if (lengths_str_arg.empty()) {
+    lengths_str_arg = bench_args.prefill_lengths;
+  }
+  std::vector<std::string> history_lengths_strs =
+      gcpp::SplitString(lengths_str_arg, ',');
+
+  size_t prefill_len = bench_args.prefill_length;
+  if (prefill_len == 0 && bench_args.generation_prefill_length > 0) {
+    prefill_len = bench_args.generation_prefill_length;
+  }
+  if (prefill_len == 0) prefill_len = 1;
+
+  std::cout << "--- Generation Benchmark (Simulated KV Cache History) ---"
+            << std::endl;
+  std::cout << "Batch Size:             " << num_queries << std::endl;
+  std::cout << "Prefill Priming Tokens: " << prefill_len << std::endl;
+  std::cout << "Decode Tokens:          " << decode_tokens << std::endl;
+  std::cout << "---------------------------------------------------------"
+            << std::endl;
+
+  for (const auto& len_str : history_lengths_strs) {
+    if (len_str.empty()) continue;
+    size_t history_len;
+    if (!gcpp::ParseSizeT(len_str, &history_len)) {
+      std::cerr << "Invalid context length: " << len_str << std::endl;
+      continue;
+    }
+
+    size_t max_seq_len = gemma.Config().max_seq_len;
+    if (history_len + prefill_len + decode_tokens > max_seq_len) {
+      size_t clamped_len = (max_seq_len > prefill_len + decode_tokens)
+                               ? max_seq_len - prefill_len - decode_tokens
+                               : 0;
+      std::cerr << "Warning: history_len=" << history_len
+                << " + prefill_len=" << prefill_len
+                << " + decode_tokens=" << decode_tokens
+                << " exceeds max_seq_len (" << max_seq_len
+                << "), clamping history_len to " << clamped_len << "."
+                << std::endl;
+      history_len = clamped_len;
+    }
+
+    std::cout << "\nSimulated History Length: " << history_len << " tokens"
+              << std::endl;
+
+    size_t total_capacity = history_len + prefill_len + decode_tokens + 32;
+
+    size_t original_seq_len = args.inference.seq_len;
+    args.inference.seq_len = total_capacity;
+
+    gcpp::RuntimeConfig gen_config = runtime_config;
+    gen_config.max_generated_tokens = decode_tokens;
+    gen_config.decode_qbatch_size = num_queries;
+    gen_config.sample_func = [](size_t, size_t, gcpp::Logits,
+                                size_t) -> gcpp::TokenAndProb {
+      // Return an arbitrary non-EOS token ID so benchmarks never terminate early.
+      return gcpp::TokenAndProb{500, 1.0f};
+    };
+
+    std::vector<gcpp::KVCache> kv_caches;
+    kv_caches.reserve(num_queries);
+    for (size_t i = 0; i < num_queries; ++i) {
+      kv_caches.emplace_back(gemma.Config(), args.inference, gen_config,
+                             ctx.allocator);
+      ZeroKVCache(kv_caches.back());
+    }
+
+    args.inference.seq_len = original_seq_len;
+
+    std::vector<int> tokens = GenerateSyntheticPrompt(gemma, prefill_len);
+
+    gcpp::TimingInfo timing_info;
+
+    size_t generated = 0;
+    size_t total_prefill_tokens = num_queries * prefill_len;
+    double start_time = hwy::platform::Now();
+    double prefill_end_time = start_time;
+
+    auto stream_token = [&generated, total_prefill_tokens,
+                         &prefill_end_time](int, float) {
+      ++generated;
+      if (generated == total_prefill_tokens) {
+        prefill_end_time = hwy::platform::Now();
+      }
+      return true;
+    };
+
+    auto batch_stream_token = [&generated, total_prefill_tokens,
+                               &prefill_end_time](size_t, size_t, int,
+                                                  float) {
+      ++generated;
+      if (generated == total_prefill_tokens) {
+        prefill_end_time = hwy::platform::Now();
+      }
+      return true;
+    };
+
+    gen_config.stream_token = stream_token;
+    gen_config.batch_stream_token = batch_stream_token;
+
+    gcpp::AllQueries all_queries(
+        tokens, /*start_pos=*/history_len, /*prefix_end=*/0,
+        hwy::Span<gcpp::KVCache>(kv_caches.data(), kv_caches.size()));
+
+    gemma.GenerateBatch(gen_config, all_queries, env, timing_info);
+
+    double end_time = hwy::platform::Now();
+    double decode_seconds = end_time - prefill_end_time;
+    size_t actual_decode_tokens =
+        generated > total_prefill_tokens ? generated - total_prefill_tokens
+                                         : 0;
+
+    std::cout << "Decode:   " << actual_decode_tokens << " tokens in "
+              << decode_seconds << "s ("
+              << (decode_seconds > 0 ? actual_decode_tokens / decode_seconds
+                                     : 0.0)
+              << " tok/s)" << std::endl;
+    std::cout << "---------------------------------------------------------"
+              << std::endl;
+  }
+
+  return 0;
+}