From c48a9aee48114a1c88f14de931ea0f94a162b8f4 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Tue, 28 Jul 2026 08:03:47 -0700 Subject: [PATCH] Helper for doing benchmarks, supports prefill , generation in both single query and multi query for generation PiperOrigin-RevId: 955262243 --- BUILD.bazel | 17 ++ evals/attention_benchmark.cc | 336 +++++++++++++++++++++++++++++++++++ 2 files changed, 353 insertions(+) create mode 100644 evals/attention_benchmark.cc diff --git a/BUILD.bazel b/BUILD.bazel index d883910f..cb93ac75 100644 --- a/BUILD.bazel +++ b/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"], diff --git a/evals/attention_benchmark.cc b/evals/attention_benchmark.cc new file mode 100644 index 00000000..a83ed024 --- /dev/null +++ b/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 + +#include +#include +#include +#include +#include +#include + +#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 { + 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 + 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 SplitString(const std::string& str, + char delimiter) { + std::vector 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 GenerateSyntheticPrompt(const gcpp::Gemma& gemma, + size_t length) { + std::vector tokens; + tokens.reserve(length); + tokens.push_back(gcpp::BOS_ID); + + int valid_token = 500; // Arbitrary non-EOS, non-control token ID fallback. + std::vector 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 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 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 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(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; +}