Repository navigation
Accumulate ARM matrix-accumulation flash attention directly into att_out. - #1076
Merged
Merged
Conversation
copybara-service
Bot
force-pushed
the
test_1012917607
branch
from
November 4, 2026 15:47
62f6522 to
6d45405
Compare
…out.
Follow-up to the generic kernel change: the matrix accumulation kernel
(TileFlashAttentionReturnExpSumsAndMaxLogitsBF16_Impl, used by the
MatrixAccumulation and MatrixAccumulationInt8 dispatches) now uses the IN/OUT
att_out directly as its online-softmax accumulator.
- TileFlashAttentionSVBlock{,BF16,Int8} take att_out's base pointer and row
stride instead of a packed C buffer. The caller always passes
kNumQueries == actual query count, and the odd row of a query pair is
already guarded, so no row >= q_count is touched.
- Remove the C region from MatrixAccumulationWorkspaceLayout, its per-call
zero fill and the final C -> att_out copy.
- HWY_DASSERT that att_out has q_count rows and qkv_dim is a multiple of 8,
since rows are now addressed directly.
- Extend the shared test setup to the matrix accumulation kernels (renaming
MixedRangesQueryType to TiledTestKernel, since it now selects kernels rather
than query types) and add
TestTiledFlashAttentionMatrixAccumulationChainedCalls, which runs the
chaining check on them.
PiperOrigin-RevId: 1013854379
copybara-service
Bot
force-pushed
the
test_1012917607
branch
from
November 4, 2026 15:59
6d45405 to
2dcef09
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 ARM matrix-accumulation flash attention directly into att_out.
Follow-up to the generic kernel change: the matrix accumulation kernel
(TileFlashAttentionReturnExpSumsAndMaxLogitsBF16_Impl, used by the
MatrixAccumulation and MatrixAccumulationInt8 dispatches) now uses the IN/OUT
att_out directly as its online-softmax accumulator.
stride instead of a packed C buffer. The caller always passes
kNumQueries == actual query count, and the odd row of a query pair is
already guarded, so no row >= q_count is touched.
zero fill and the final C -> att_out copy.
since rows are now addressed directly.
MixedRangesQueryType to TiledTestKernel, since it now selects kernels rather
than query types) and add
TestTiledFlashAttentionMatrixAccumulationChainedCalls, which runs the
chaining check on them.