From e7bc886aef1deb052e015fe6d62fcd3aaff71372 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Fri, 28 Aug 2026 05:31:36 -0700 Subject: [PATCH] Add ability to benchmark prefill in attention benchmark binary PiperOrigin-RevId: 972529744 --- evals/attention_benchmark.cc | 119 ++++++++++++++++++++--------------- 1 file changed, 67 insertions(+), 52 deletions(-) diff --git a/evals/attention_benchmark.cc b/evals/attention_benchmark.cc index a83ed024..f79bf513 100644 --- a/evals/attention_benchmark.cc +++ b/evals/attention_benchmark.cc @@ -36,6 +36,9 @@ // ============================================================================= // Supported Benchmark Flags // ============================================================================= +// --benchmark_mode: Benchmark mode: 'both' (benchmarks prefill and +// decode) or 'decode' (benchmarks decode with simulated +// KV cache history). Default: 'both'. // --context_lengths: Comma-separated KV cache history lengths to // benchmark (default: '1024,8192,32768'; aliases: // --generation_context_lengths, --prompt_lengths) @@ -77,6 +80,7 @@ class AttentionBenchmarkArgs : public ArgsBase { std::string generation_context_lengths; std::string prompt_lengths; std::string prefill_lengths; + std::string benchmark_mode; size_t prefill_length; size_t generation_prefill_length; size_t benchmark_tokens; @@ -84,6 +88,10 @@ class AttentionBenchmarkArgs : public ArgsBase { template void ForEach(const Visitor& visitor) { + visitor(benchmark_mode, "benchmark_mode", std::string("both"), + "Benchmark mode: 'both' (prefill and decode) or 'decode' (decode " + "with simulated KV cache history)", + 2); visitor(context_lengths, "context_lengths", std::string("1024,8192,32768"), "Comma-separated list of simulated KV cache history lengths to " @@ -223,40 +231,49 @@ int main(int argc, char** argv) { } if (prefill_len == 0) prefill_len = 1; - std::cout << "--- Generation Benchmark (Simulated KV Cache History) ---" - << std::endl; + const bool is_only_decode = (bench_args.benchmark_mode == "decode"); + + std::cout << "--- Benchmark: " << (is_only_decode ? "decode" : "both") + << " ---" << std::endl; std::cout << "Batch Size: " << num_queries << std::endl; - std::cout << "Prefill Priming Tokens: " << prefill_len << std::endl; + if (is_only_decode) { + 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)) { + size_t target_len; + if (!gcpp::ParseSizeT(len_str, &target_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; + const size_t max_seq_len = gemma.Config().max_seq_len; + const size_t priming_tokens = is_only_decode ? prefill_len : 0; + if (target_len + priming_tokens + decode_tokens > max_seq_len) { + const size_t max_allowed = + (max_seq_len > priming_tokens + decode_tokens) + ? max_seq_len - priming_tokens - decode_tokens + : 0; + std::cerr << "Warning: target_len=" << target_len + << " exceeds max_seq_len (" << max_seq_len << "), clamping to " + << max_allowed << "\n"; + target_len = max_allowed; } - std::cout << "\nSimulated History Length: " << history_len << " tokens" - << std::endl; + const size_t history_len = is_only_decode ? target_len : 0; + const size_t prompt_tokens = is_only_decode ? prefill_len : target_len; + + std::cout << "\n" + << (is_only_decode ? "Simulated History Length: " + : "Context / Prefill Length: ") + << target_len << " tokens\n"; - size_t total_capacity = history_len + prefill_len + decode_tokens + 32; + size_t total_capacity = + std::min(max_seq_len, history_len + prompt_tokens + decode_tokens + 32); size_t original_seq_len = args.inference.seq_len; args.inference.seq_len = total_capacity; @@ -280,56 +297,54 @@ int main(int argc, char** argv) { args.inference.seq_len = original_seq_len; - std::vector tokens = GenerateSyntheticPrompt(gemma, prefill_len); + std::vector tokens = GenerateSyntheticPrompt(gemma, prompt_tokens); gcpp::TimingInfo timing_info; size_t generated = 0; - size_t total_prefill_tokens = num_queries * prefill_len; - double start_time = hwy::platform::Now(); + const size_t total_prefill_tokens = num_queries * prompt_tokens; + const 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) { + auto on_token = [&]() { + 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 = [&](int, float) { return on_token(); }; + gen_config.batch_stream_token = [&](size_t, size_t, int, float) { + return on_token(); }; - 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; + const double end_time = hwy::platform::Now(); + const double prefill_seconds = prefill_end_time - start_time; + const double decode_seconds = end_time - prefill_end_time; + const size_t actual_decode_tokens = + generated > total_prefill_tokens ? generated - total_prefill_tokens : 0; + + if (!is_only_decode) { + std::cout << "Prefill: " << total_prefill_tokens << " tokens in " + << prefill_seconds << "s (" + << (prefill_seconds > 0 ? total_prefill_tokens / prefill_seconds + : 0.0) + << " tok/s); TTFT: " << static_cast(prefill_seconds * 1E3) + << " ms\n"; + } + if (decode_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)\n"; + } + std::cout << "---------------------------------------------------------\n"; } return 0;