From 2c9669bb1236af604041c40480ad257ef72c2ed8 Mon Sep 17 00:00:00 2001 From: Dani Ferreira Franco Moura Date: Wed, 14 Oct 2026 06:52:21 -0700 Subject: [PATCH] Document that last_pos_per_query is inclusive in flash attention. The masks in SingleFlashAttentionRowVector and DoubleFlashAttentionRowVector keep keys in [pos, last_pos] (mask_size = last_pos + 1 - pos), and all callers pass the index of the last attended key (e.g. tiled_attention.cc uses the query's own position, vit.cc uses num_tokens - 1). The ComputeQKMacroTile doc block said "(exclusive)" and the public header comment was unclear. Comment-only change, no behavior change. PiperOrigin-RevId: 999337547 --- gemma/flash_attention.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index edba20df..eb6d5233 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -1094,8 +1094,8 @@ static HWY_INLINE void QDotKTilexUpTo8TransposedKDoubleWidthBF16( // q_T_in_groups_up_to_4 - Span of float* All except last float* // should have (qkv_dim, 4) Last one can have any size up to 4. // q_scales - Span of float of shape (q_count,) used for Queries in int16 format -// start_pos_per_query - start position in kv to start attention from () -// last_pos_per_query - last position in kv to attend to (exclusive) +// start_pos_per_query - first position in kv to attend to (inclusive) +// last_pos_per_query - last position in kv to attend to (inclusive) // queries_per_timestep - how many queries begin/end on the same timestep // attention_shape - see struct definition for more details. // att_cap - soft cap on attention logits