Use unroll_count(4) for the NAX attention Q@K.T loop#3843
Open
wyanzhao wants to merge 1 commit into
Open
Conversation
Fully unrolling the head-dim loop makes the compiler hoist all TD K-tile
loads to the top of the block, ahead of the mma chain. Unrolling by 4
instead lets it interleave each K load with the running matmul, which
measures faster at head_dim 128 and is a no-op at head_dim 64.
M5 Max, 16 Q / 2 KV heads, bf16, causal, 8192^2 (median of 12):
head_dim 128: 5.70 -> 5.08 ms (48.2 -> 54.1 TF, +12%)
head_dim 64: 2.45 -> 2.44 ms (unchanged; TD == 4 is still a full
unroll at count 4)
No numerical change (bitwise-identical accumulation order); 16/16
fast-sdpa tests pass.
wyanzhao
marked this pull request as ready for review
July 13, 2026 20:54
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.
Proposed changes
Change the NAX attention
Q @ K.Thead-dimension loop from full unrollingto
#pragma clang loop unroll_count(4).Note
This is a scheduling-only change. The accumulation order and numerical
results are unchanged.
Motivation
With full unrolling, the compiler hoists all
TDK-tile loads to the startof the block, ahead of the running MMA chain. Unrolling by four instead lets
the compiler interleave K-tile loads with matrix accumulation.
unroll_count(4)head_dim == 128(TD == 8)head_dim == 64(TD == 4)Performance
M5 Max, 16 Q / 2 KV heads, bf16, causal,
8192 x 8192; median of 12 runs:1285.70 ms/48.2 TF5.08 ms/54.1 TF+12%throughput642.45 ms2.44 msScope and compatibility
wn=1NAX QK loop.head_dim == 256,wn=2D-split path.head_dim == 256PR; the two changes touchdifferent regions of
steel_attention_nax.hand can merge in either order.Validation
order is unchanged.
test_fast_sdpa.pyresult: 16 tests, 1 platform-dependentskip, otherwise passing.
head_dim == 256PR completes without conflict.Checklist
pre-commit run --all-files.