Adds Intel AMX TMUL hardware acceleration for FlashAttention using BF16 - #1039
Merged
Merged
Conversation
copybara-service
Bot
force-pushed
the
test_977877681
branch
from
September 22, 2026 13:49
c22f5f3 to
9e800ac
Compare
matrix multiplication: - Implements `TileFlashAttentionAMX_BF16_Impl` on top of Highway's AMX-BF16 wrappers (`Tile64BLoad` / `Tile64BStore` / `Tile64BMatMul`), so the compiler allocates tile registers and emits the tile configuration rather than this kernel hand-rolling `_tile_*` intrinsics and a packed `AMXTileConfig`. - Gates the kernel on `hwy::HaveTile64BMatMulBF16()`, which additionally validates XSAVE/OSXSAVE and XCR0 before requesting `XFEATURE_XTILEDATA`. - Integrates `AttentionImpl::kFlashAMX` into gemma_cpp tiled attention dispatch and model forward passes. - Adds comprehensive unit tests (`flash_attention_test`, `tiled_attention_test`) and benchmarks. About x4 speedup on supported hardware PiperOrigin-RevId: 985969928
copybara-service
Bot
force-pushed
the
test_977877681
branch
from
September 22, 2026 14:02
9e800ac to
94280dd
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
1 participant
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.
Adds Intel AMX TMUL hardware acceleration for FlashAttention using BF16
matrix multiplication:
TileFlashAttentionAMX_BF16_Implon top of Highway's AMX-BF16wrappers (
Tile64BLoad/Tile64BStore/Tile64BMatMul), so the compilerallocates tile registers and emits the tile configuration rather than this
kernel hand-rolling
_tile_*intrinsics and a packedAMXTileConfig.hwy::HaveTile64BMatMulBF16(), which additionallyvalidates XSAVE/OSXSAVE and XCR0 before requesting
XFEATURE_XTILEDATA.AttentionImpl::kFlashAMXinto gemma_cpp tiled attention dispatch and model forward passes.flash_attention_test,tiled_attention_test) and benchmarks.About x4 speedup on supported hardware