Skip to content

Document that last_pos_per_query is inclusive in flash attention. - #1058

Merged
copybara-service[bot] merged 1 commit into
devfrom
test_999310118
Oct 14, 2026
Merged

copybara-service[bot] merged 1 commit into
devfrom
test_999310118

Conversation

@copybara-service

Copy link
Copy Markdown

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.

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
@copybara-service
copybara-service Bot merged commit 2c9669b into dev Oct 14, 2026
6 checks passed
@copybara-service
copybara-service Bot deleted the test_999310118 branch October 14, 2026 13:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant