Skip to content

Accumulate generic tile flash attention directly into att_out. - #1065

Merged
copybara-service[bot] merged 1 commit into
devfrom
test_1003900950
Nov 3, 2026
Merged

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

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Oct 21, 2026 •

Copy link
Copy Markdown

Accumulate generic tile flash attention directly into att_out.

att_out is now an IN/OUT parameter: the generic
TileFlashAttentionReturnExpSumsAndMaxLogits kernel uses it directly as the
online-softmax accumulator instead of a separate C accumulator buffer.

  • Remove the C accumulator region from TileFlashAttentionWorkspaceLayout
    (saves padded_q_count * qkv_dim * 4 bytes of per-worker workspace).
  • Drop the per-call zero fill of C and the final C -> att_out copy.
  • Address att_out via Row()/Stride() so padded layouts work.
  • Document att_out as in-out, like exp_denominator_sums and max_logits: it
    must be zero-initialized before the first call, and its rows must be
    aligned to the SIMD vector size (all current callers already satisfy both).
    Fix the MixedRanges test and the attention benchmark, which passed
    uninitialized att_out.
  • Add TestTiledFlashAttentionChainedCalls with a kernel-agnostic helper,
    CheckChainedCallsMatchSingleCall, verifying that chaining calls over
    sub-ranges of the KV (carrying att_out, exp_sums and max_logits) matches a
    single call over the full range, including remainder query groups. The
    MixedRanges and chaining tests share a TiledFlashAttentionTestSetup.

att_out is now an IN/OUT parameter: the generic
TileFlashAttentionReturnExpSumsAndMaxLogits kernel uses it directly as the
online-softmax accumulator instead of a separate C accumulator buffer.

- Remove the C accumulator region from TileFlashAttentionWorkspaceLayout
  (saves padded_q_count * qkv_dim * 4 bytes of per-worker workspace).
- Drop the per-call zero fill of C and the final C -> att_out copy.
- Address att_out via Row()/Stride() so padded layouts work.
- Document att_out as in-out, like exp_denominator_sums and max_logits: it
  must be zero-initialized before the first call, and its rows must be
  aligned to the SIMD vector size (all current callers already satisfy both).
  Fix the MixedRanges test and the attention benchmark, which passed
  uninitialized att_out.
- Add TestTiledFlashAttentionChainedCalls with a kernel-agnostic helper,
  CheckChainedCallsMatchSingleCall, verifying that chaining calls over
  sub-ranges of the KV (carrying att_out, exp_sums and max_logits) matches a
  single call over the full range, including remainder query groups. The
  MixedRanges and chaining tests share a TiledFlashAttentionTestSetup.

PiperOrigin-RevId: 1012927185
@copybara-service
copybara-service Bot merged commit fef595d into dev Nov 3, 2026
4 of 5 checks passed
@copybara-service
copybara-service Bot deleted the test_1003900950 branch November 3, 2026 13:17
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