Repository navigation
Accumulate AMX flash attention directly into att_out. - #1077
Merged
Merged
Conversation
copybara-service
Bot
force-pushed
the
test_1012919692
branch
from
November 4, 2026 16:27
8ee2d1b to
8ccdbad
Compare
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
Bot
force-pushed
the
test_1012919692
branch
from
November 4, 2026 16:41
8ccdbad to
1809753
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
StoreAccumulatorsToOutput) with row-major access to att_out through its
MatPtrT (AccumRow helper).
it accumulates into att_out directly. Remove the C region from
AmxInt8WorkspaceLayout, its zero fill and the final copy.
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.
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.