metal: wide-tile indexer prefill scorer over half-packed Q/K (bit-exact, +4.7% prefill at 64k) - #830
Open
adriangalilea wants to merge 1 commit into
Open
Conversation
The pre-M5 prefill scorer (kernel_dsv4_indexer_scores_tiled) re-reads the full f32 Q tile from device memory once per head per 32-row comp tile -- about 1 KB of the ~1.1 KB it moves per scored (token, comp) pair, and the dominant term of the long-context prefill slope (the score stage is 56% of the growing attention cost at 64k context on M3 Ultra + Flash 0731 MXFP4). Split it: two f32->f16 pack dispatches convert Q and the indexer K rows once per call with exactly the half(float) rounding the kernel applied during staging, and a TN=64 variant (kernel_dsv4_indexer_scores_tiled2_f16, 256 threads, 8 simdgroups, 20.5 KB threadgroup memory) reads the packed halves. Half the comp tiles means half the Q re-reads, and each re-read moves half the bytes. Per (token, comp) pair the reduction is unchanged -- heads ascending, sixteen 8x8 simdgroup MACs per head in the same order, relu(dot)*w accumulated in float -- so scores are bit-identical. Gated on batched prefill (n_tokens >= 32), off in --quality; DS4_METAL_DISABLE_INDEXER_SCORES_TILED2 rolls back, read per call so the variant bench can toggle it in-process.
This was referenced Aug 18, 2026
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.
The pre-M5 prefill indexer scorer (
kernel_dsv4_indexer_scores_tiled) re-reads the full f32 Q tile from device memory once per head per 32-row comp tile — about 1 KB of the ~1.1 KB it moves per scored (token, comp) pair. The kernel's own comment notes the score matrix "dominates the prefill slope"; this PR attacks the traffic term without touching the arithmetic.Split it in two: a pair of
f32→f16pack dispatches convert Q and the indexer K rows once per call, with exactly thehalf(float)rounding the kernel applied during staging (the rounding moves, the values don't), and a TN=64 variant (kernel_dsv4_indexer_scores_tiled2_f16, 256 threads, 8 simdgroups, 20.5 KB threadgroup memory) reads the packed halves. Half the comp tiles → half the Q re-reads, each moving half the bytes. Per (token, comp) pair the reduction is unchanged — heads ascending, sixteen 8×8 simdgroup MACs per head in the same order,relu(dot)*waccumulated in float — so scores are bit-identical. Gated on batched prefill (n_tokens >= 32), off in--quality;DS4_METAL_DISABLE_INDEXER_SCORES_TILED2rolls back, read per call so the variant bench can toggle it in-process.Measured on Apple M3 Ultra (60-core GPU) 512 GB, DeepSeek V4 Flash 0731 native-MXFP4 GGUF (156 GB), base 84cc882:
Exactness:
metal_prefill_variant_bench --candidate-env DS4_METAL_DISABLE_INDEXER_SCORES_TILED2 --prefix-tokens 8192 --repeats 2→ 8/8 runs bit-exact,exact_floats=1034240, candidate (rollback) −0.62%. Full-vocab logits at frontiers 2048..8192 byte-identical to a build without the change (ds4-bench --dump-frontier-logits-dirbyte-compare across binaries). Harness control-vs-control calibrated bit-exact on this box (see my comment in #793).Speed — the CONTRIBUTING.md sweep (
--ctx-start 2048 --ctx-max 65536 --step-incr 2048 --gen-tokens 128, ABBA pairs, quiet box), delta = this PR vs 84cc882:Monotone in context (it is a traffic fix for a term that grows with n_comp); generation deltas are noise. Single cold syncs at
--prefill-chunk 4096: 16k 590→599, 32k 550→563, 64k 484→504 t/s. TheDS4_METAL_INDEXER_STAGE_PROFILEsplit shows the score stage at 64k dropping 25.9 s → 20.6 s per full prefill (−21%).make testtargets green incl../ds4_test --server; built with zero warnings. Running in production on the machine above. Scratch cost: two shared buffers,n_tokens×n_head×128andn_comp×128halves (~64 + ~67 MB at 1M ctx / 4096 chunk).Related: this is the Metal analogue of the mxf4 indexer-scorer direction CUDA already took (56ec892, 08fecd9) but kept bit-exact; a follow-up porting the mxf4 encoding itself (matching CUDA's default, not bit-exact vs today's Metal) would cut the remaining byte term a further ~8x.