From 7a8c621ffb306540ab6bfafb32301f2ce7dd6302 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Fri, 17 Jul 2026 01:54:51 -0700 Subject: [PATCH] Pass runtime_config_ to the KVCache constructor within GemmaEnv in benchmark_helper.cc. This ensures that kv_cache_type and attention_impl settings are properly propagated when using non-default attention implementations such as flash_transposed_qs_bf16. PiperOrigin-RevId: 949438536 --- evals/benchmark_helper.cc | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/evals/benchmark_helper.cc b/evals/benchmark_helper.cc index 71869d14..f9a846a7 100644 --- a/evals/benchmark_helper.cc +++ b/evals/benchmark_helper.cc @@ -42,8 +42,6 @@ GemmaEnv::GemmaEnv(const GemmaArgs& args) env_(ctx_), gemma_(args, ctx_) { const ModelConfig& config = gemma_.Config(); - // Only allocate one for starters because GenerateBatch might not be called. - kv_caches_.push_back(KVCache(config, args.inference, ctx_.allocator)); if (args.inference.verbosity >= 2) { ShowConfig(args, config, gemma_.WeightReadMode(), ctx_); @@ -57,6 +55,10 @@ GemmaEnv::GemmaEnv(const GemmaArgs& args) .verbosity = args.inference.verbosity, }; args.inference.CopyTo(runtime_config_); + + // Only allocate one for starters because GenerateBatch might not be called. + kv_caches_.push_back( + KVCache(config, args.inference, runtime_config_, ctx_.allocator)); } QueryResult GemmaEnv::QueryModel(const std::vector& tokens) { @@ -130,8 +132,8 @@ QueryResultAndMetrics GemmaEnv::BatchQueryModelWithMetrics( // Ensure we have at least one KVCache per query. while (kv_caches_.size() < num_queries) { - kv_caches_.push_back( - KVCache(gemma_.Config(), gemma_.Inference(), ctx_.allocator)); + kv_caches_.push_back(KVCache(gemma_.Config(), gemma_.Inference(), + runtime_config_, ctx_.allocator)); } const hwy::Span kv_caches(&kv_caches_[0], num_queries);