[gilby]

What this changes

An opt-in WMMA kernel for Flash Next's sparse attention on prompt chunks on gfx11 (--policy sparse=wmma, default scalar): 4 query rows x the kv head's 12 query heads as 48 WMMA columns over the union of those rows' selected blocks, per-row block masks and causal ends, fp32 online softmax, the existing parts/merge layout. Only prompt chunks take it (o.prefill); MTP verify windows, lane rounds and decode keep tf_fn_sparse_parts12. Fixes #1597.

This changes the prompt's attention sums (Q, K, P and V rounded to the cache's 16-bit type in WMMA fragments with fp32 accumulation, and a different key partition and order), so it needs the arithmetic gate agreed in #1597 before it could become a default; it stays opt-in here.

Receipt

• Environment: TensorFold 7c068895 plus this change, tensorfold-native 1.1.8 built with zig build native-hip -Dhip-arch=gfx1151, Ubuntu archive HIP 7.1 (clang 21), AMD Ryzen AI MAX+ 395 with Radeon 8060S gfx1151 (128 GB unified memory). Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6, MTP 3 drafts, --parallel 1 --context 131072 --prompt-cache-gib 0, base and candidate on the same machine.
• Exactness with sparse=wmma: tools/bench_concurrent.py --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial every cell alone 3/3 and 12/12 equal and serial (against "draft": false) equal, 0 failed, also with prefill_step=4096. Resumed equals fresh: a two-turn chat answered by a server with --prompt-cache-gib 6 (turn 2 resumed at 10,944 and 12,224 tokens) gives the same reply token_sha256 as a fresh process with the cache off, under sparse=wmma and under the scalar kernel; resume points sit on the prompt cache's 64-row grid, which is also the attention slice, so every 4-row tile starts at the same position in both runs. Against the scalar kernel the outputs differ (the sums change): the 16-token replies to both cold prompts are identical, the two 256-token greedy replies diverge partway.
• Kernel accuracy: against an fp64 reference on served shapes (24/2 heads, 256 dims, 64-row slice; 6/1 heads, 17-row tail; a dense row), worst absolute error 2.2e-3 (scalar kernel 2.6e-7).
• Decode speed, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, median tok/s code / chat: base 49.8 / 43.8, sparse=wmma 47.1 / 40.1. Decode runs the same scalar kernel in both; the replies themselves differ (prompt sums change), and with them MTP acceptance.
• Prompt speed, tools/bench_cache.py --require-cold --tokens 16 on public 7,377- and 29,364-token arrays, first cold request ttft_s: base 48.8 / 111.0 s; sparse=wmma 40.4 / 74.6 s; sparse=wmma,prefill_step=4096 39.2 / 66.9 s. The kernel's own time on a 7.9K-token prompt drops from 9.81 s to 0.73 s (rocprofv3).
• Tests run: zig build hip-kernel-contract on gfx1151: SUMMARY gfx1151 50 kernels: 50 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems (base: 49 kernels, the same 126 live-versus-recorded register differences for existing kernels under this machine's compiler); the new entries: non-finite (a NaN or Inf a row did not select does not reach it; a selected NaN propagates), repeat, 7 shapes within the fp64 bound (worst 0.021 of it), sums, and live limits matching their records. zig build test -Dhip=true --summary all 466/466 steps, 0 failed. zig build hip-kernel-test: the new sparse prompt test passes; the run then stops at the same Nemotron test as the base build on this machine. python3 tools/zig/lean_check.py 425 on base and change. git diff --check clean.
• Not run: hip-kernel-contract-limits stops on this machine's compiler at tiles/dot2_tiles.hip for gfx1100 (branch size exceeds simm16, same on the base, see HIP: dot2_tiles.hip fails to compile for gfx1100 with Ubuntu clang 21.1.8 ("branch size exceeds simm16") #1521); records for gfx1030, gfx1100 and gfx1201 come from each chip's compiled code object, and those chips keep the scalar kernel at runtime (the policy is ANDed with gfx11 WMMA support). Resume points off the 64-row grid (explicit request chunks). Other checkpoints and AMD chips; Metal and CUDA.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
• No precision traded for speed. If the bits change, the description says which sums change and why. (Opt-in; the prompt sums change as described above, gate to agree in HIP Flash Next: WMMA sparse attention for prompt chunks on gfx11 #1597.)
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks for this. A WMMA tile for sparse prefill is a big piece of work, and keeping it opt-in, with contract records and an fp64 test, is the right shape.

The main question is exactness. The tile rounds q*scale and the softmax probabilities to 16 bits before the products, the contract spec marks it .exact = false, and the test allows 5e-3 against fp64 where the scalar kernel is held to 1e-5. So turning on sparse=wmma will change output tokens. We hold the engine to exact paths, and an opt-in that changes tokens needs a strong case and proof. Could you try keeping q and P in fp32 for the products, or otherwise matching the scalar kernel's bits, even if part of the speedup goes away? Either way, please attach before/after prefill numbers, the device receipts for the two contract records and the kernel test, and how often tokens differ with it on.

Two smaller points:

1. The launcher assumes a compress ratio of 4 (kFxRatio). Please check the model's ratio and refuse otherwise.
2. The Choices comment ("two give the same bits; every one defaults on") is no longer true with this key, so it needs updating.

Also, #1614 edits the same lines in hip_records.zig and contract_main.zig, so whichever lands second will need a small rebase.

[gilby]

Thanks for the careful review. Short version: q and P are now fp32 in the products, and the tile is at the scalar kernel's fp32 error against fp64. It still does not match the scalar kernel's bits, and output tokens still differ, so I can't make the exactness case you asked for. Details below so you can decide whether it is worth keeping as an opt-in or should be closed.

What changed

• 435d76e9d: q*scale and P are each carried as three 16-bit parts (hi, mid, lo) whose fp32 sum is the value, with one matrix op per part and fp32 accumulation; K and V stay in the cache type. Three Q fragment sets spilled at 256 threads (256 VGPRs, 348 B scratch a lane), so a block is now 512 threads, four waves splitting the head dimension of each score tile: 189 VGPRs, no scratch, 39,936 B LDS on gfx1151. The 16-bit variant is gone, so there is one option.
• The kernel test now holds the tile to 1e-5 against fp64, the scalar kernel's bound (it was 5e-3; the old 16-bit tile fails the new bound), and a new deep case (2,048 selected keys a row) compares tile to scalar directly at 1e-5. The contract spec drops the 16-bit allowance and holds the tile to the fp32 depth bound; it stays .exact = false.
• 1bbc1075a: Engine.load refuses sparse=wmma when indexer_compress_ratio is not 4 (sparse=wmma needs indexer_compress_ratio 4; this model has N, error.SparseWmmaRatio); the scalar path takes any ratio. A host test reads the fixture config (4, accepted) and checks 2 and 8 are refused with wmma while the scalar path takes 2. The Choices comment now reads "fuse and wide give the same bits and default on; sparse_wmma is opt-in".

Why the bits still differ

The scalar kernel splits a row's keys into parts by position in its key list, walks 32-key tiles with sequential fp32 FMAs and keeps a per-row online max. The tile splits parts by cache block, walks 16-key union tiles and accumulates in the matrix cores' order. Those are different fp32 sums of the same terms, so they round differently in the last bits, and a 1e-7 attention difference is enough to flip a later bf16 rounding and then a greedy token. I don't see a way to reproduce the scalar order with matrix-core accumulation without giving up the tile.

Receipts (gfx1151, bf16 cache)

Max abs error, kernel test shapes:

shape	scalar vs fp64	old 16-bit tile vs fp64	fp32 tile vs fp64	fp32 tile vs scalar
24 heads x 64 rows	2.6e-7	2.2e-3	3.1e-7	5.1e-7
6 heads x 17 rows	1.9e-7	2.0e-3	2.7e-7	3.3e-7
2,048 keys a row				1.65e-6

Kernel time for a 64-row slice with 2,048 selected keys a row: scalar 8.71 ms, old 16-bit tile 0.98 ms, fp32 tile 1.18 ms.

Served, TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP, --parallel 1 --context 131072 --policy prefill_step=4096 --prompt-cache-gib 0, tools/bench_cache.py --require-cold, steady cold ttft:

prompt	scalar (default)	sparse=wmma (fp32 tile)
7,377 tokens	21.41 s	13.43 s
29,364 tokens	94.44 s	58.11 s

How often tokens differ from the scalar path, T=0:

• tools/bench_openai.py 256-token replies (fibonacci-raw, gpu-chat-no-think): 0 of 2 identical. token_sha256 scalar vs tile: beb57dfacf5e vs 6fd63835d6af and e5fa3204ab36 vs 6136d98947ff.
• 10 more prompts (5 chat, 5 completion), 256-token replies: 1 of 10 identical; the first difference falls between character 1 and about 890 of replies of roughly 950-1,400 characters.
• The two cold prompts' 16-token replies: 2 of 2 identical.
• With the tile on, replies are deterministic (one hash per prompt over 15 repeats), and bench_concurrent --alone --serial is equal in every cell, so drafted equals plain and concurrent equals solo within the option.
• Decode (the tile only runs on prompt chunks, so this follows the different replies and their MTP acceptance): fibonacci-raw 49.8 vs 45.8 tok/s, gpu-chat-no-think 43.8 vs 44.1.

Gates on this branch: hip-kernel-test 105/105; hip-kernel-contract SUMMARY gfx1151 50 kernels: 50 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems (main: 49 kernels, the same 126 register differences in other kernels); both sparse records pass shapes (worst 0.004 of the bound), nonfinite, repeat (5 runs byte-identical), sums and live limits. zig build test -Dhip=true 466/466 steps; lean 425 as on main; zig fmt --check and git diff --check clean. The gfx1030, gfx1100 and gfx1201 records come from hipcc -Rpass-analysis=kernel-resource-usage: hip-kernel-contract-limits stops at tiles/dot2_tiles.hip for gfx1100 on this machine's compiler, same on main. The f16-cache instance compiles and passes the contract but has no served model on this chip.

So the option is now "fp32-accurate, 1.6x faster cold prefill, not bit-identical, tokens differ". If that is not enough for an opt-in under the exactness rule, I'm fine closing this; it was worth knowing where the line is.