Skip to content

Accumulate AMX flash attention directly into att_out. - #1077

Merged
copybara-service[bot] merged 1 commit into
devfrom
test_1012919692
Nov 4, 2026
Merged

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

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

Accumulate AMX flash attention directly into att_out.

Follow-up to the generic and matrix accumulation kernel changes: the AMX-BF16
and AMX-INT8 kernels now use the IN/OUT att_out directly as their
online-softmax accumulator.

  • Replace the blocked accumulator layout (GetAccumTilePtr /
    StoreAccumulatorsToOutput) with row-major access to att_out through its
    MatPtrT (AccumRow helper).
  • AMX-INT8: accumulator rows are only accessed per row for rows < q_count, so
    it accumulates into att_out directly. Remove the C region from
    AmxInt8WorkspaceLayout, its zero fill and the final copy.
  • AMX-BF16: remove the per-call heap allocation and zero fill of C. Full
    32-query blocks Tile64BLoad/Store straight from att_out with its row
    stride. AMX tiles span all 32 rows of a block, so a partial last block uses
    a MatPtrT view over a 32 x qkv_dim scratch vector (allocated only when
    there is a partial block), copied in from and back out to att_out.
  • Document att_out as accumulated in place by all kernels.
  • Extend the shared test setup to the AMX kernels and add
    TestTiledFlashAttentionAMXChainedCalls with 64, 40 and 15 queries (full
    blocks, full + partial, partial only), via a new query_counts parameter of
    RunTiledFlashAttentionChainingTest. It skips visibly (GTEST_SKIP) on hosts
    without AMX instead of silently passing.

Follow-up to the generic and matrix accumulation kernel changes: the AMX-BF16
and AMX-INT8 kernels now use the IN/OUT att_out directly as their
online-softmax accumulator.

- Replace the blocked accumulator layout (GetAccumTilePtr /
  StoreAccumulatorsToOutput) with row-major access to att_out through its
  MatPtrT (AccumRow helper).
- AMX-INT8: accumulator rows are only accessed per row for rows < q_count, so
  it accumulates into att_out directly. Remove the C region from
  AmxInt8WorkspaceLayout, its zero fill and the final copy.
- AMX-BF16: remove the per-call heap allocation and zero fill of C. Full
  32-query blocks Tile64BLoad/Store straight from att_out with its row
  stride. AMX tiles span all 32 rows of a block, so a partial last block uses
  a MatPtrT view over a 32 x qkv_dim scratch vector (allocated only when
  there is a partial block), copied in from and back out to att_out.
- Document att_out as accumulated in place by all kernels.
- Extend the shared test setup to the AMX kernels and add
  TestTiledFlashAttentionAMXChainedCalls with 64, 40 and 15 queries (full
  blocks, full + partial, partial only), via a new query_counts parameter of
  RunTiledFlashAttentionChainingTest. It skips visibly (GTEST_SKIP) on hosts
  without AMX instead of silently passing.

PiperOrigin-RevId: 1013882475
@copybara-service
copybara-service Bot merged commit 1809753 into dev Nov 4, 2026
1 check passed
@copybara-service
copybara-service Bot deleted the test_1012919692 branch November 4, 2026 16:41
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