perf(metal): widen four-row flash KV tiles - #89
Merged
mxaddict merged 6 commits intoJul 30, 2026
Conversation
Add a 128-position KV specialization for the four-row, hd128 flash attention shape. Runtime dispatch prefers it when the pipeline fits and retains the existing 64-position specialization as fallback.
Keep the C128 shape policy inside the runtime-used resolver and update the softmax ownership comment for both KV chunk widths.
Probe the C128 pipeline only for its four-row target shape and compile the standard flash pipeline only when C128 cannot run. This avoids paying for both specializations during pipeline warmup.
Exercise four-row C128 flash attention across a non-aligned second KV chunk and with a sliding-window lower bound inside the first chunk.
digitaloten
marked this pull request as ready for review
July 28, 2026 11:31
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
2 participants
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.
Summary
Verification
cargo fmt --all -- --checkcargo test -p infr-metal attention_flash128_policy_is_exactcargo test -p infr-metal --test kernel_namescargo check -p infr-metal --all-targets --lockedcommon.metal+attention.metal; AIR exportsattnflash2_c128_f16kv_hd128Performance evidence
Paired same-runner A/B at
rows=4,n_head=16,n_kv=8,head_dim=128,kv_len=4096,pos=4092, f16 KV, causal mask. Each process used deterministic inputs, at least 50 warmups and one second of GPU warmup, then 10 calibrated 2–5 second blocks; block 0 was discarded. C128/C64 order alternated across fivemacos-15runners. The metric is one-op synchronousMetalBackend::executewall latency, not isolated kernel time.Runner-level geometric-mean ratio: 0.879009 (12.099% faster). Runner-level bootstrap 95% CI: [0.851989, 0.911652]. The predeclared win criterion was an upper bound below
0.95; this passes.Evidence run: https://github.com/digitaloten/infr/actions/runs/30352605155