From b94f4468a6786ba5a2dc66d4e783fb36527f63d3 Mon Sep 17 00:00:00 2001
From: Krzysztof Rymski <krymski@google.com>
Date: Wed, 15 Jul 2026 07:59:33 -0700
Subject: [PATCH] Adjust work spliting to subtasks, about 1.5x faster prefill
 on zen5 machines, and +- neutral on older ones. Importantly to fully utilize
 this it's important to use high prefill token_batch like 4k

PiperOrigin-RevId: 948338059
---
 gemma/tiled_attention.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc
index 22e20e0c..875565ac 100644
--- a/gemma/tiled_attention.cc
+++ b/gemma/tiled_attention.cc
@@ -605,7 +605,7 @@ void LocalAttentionForAllHeadsTokensAndBatch(
     const size_t layer_idx, const LayerWeightsPtrs& layer,
     AttentionActivationsPtrs& activations, QBatch& qbatch,
     ThreadingContext& ctx) {
-  constexpr size_t kQueriesPerSubtask = 32;
+  constexpr size_t kQueriesPerSubtask = 128;
   const size_t heads_per_kv_head =
       layer.layer_config.heads / layer.layer_config.kv_heads;
   const hwy::Divisor div_heads_per_kv_head(heads_per_kv_head);