[gilby]

What this changes

A routed prompt matrix-core tile of 64 rows an item (affine_wmma_item<BITS, 4>, all 8 waves across the 128 columns, the 128-row tile's staging and fold unchanged, so rows are byte-identical), selected through a new gfx1151.zon (gfx1100's rows plus mlx.routed.prefill.item64 at cost 9 for routed plans of up to 12 pairs an item, using #1607's pairs-per-item bound). Fixes #1601.

Receipt

• Environment: TensorFold main f350cd3e (with HIP: bound the routed prompt's K-parallel tile by pairs per item, and give it gfx1201's short prompts #1607) 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). Checkpoints TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6 and an oMLX oQ4e quantization of Qwen3.8-Flash-Next, MTP 3 drafts, --parallel 1 --context 131072 --prompt-cache-gib 0 --policy prefill_step=4096, main and candidate on the same machine.
• Exactness: token_sha256 of tools/bench_openai.py, tools/bench_concurrent.py and both tools/bench_cache.py prompts identical to main on both checkpoints; bench_concurrent --alone --serial 0 failed, 0 unmeasured, every cell equal.
• Kernel timing, zig build hip-affine-tune on gfx1151 at Flash Next's served routed plans (512 experts plus the shared one, 11 slots), us, 128-row tile -> 64-row tile: gate/up (n 1280, k 2560) 256 tokens 10,987 -> 9,741, 512 tokens 11,576 -> 9,932, 1,024 tokens 21,215 -> 10,629, 2,048 tokens 25,580 -> 12,303; down (n 2560, k 640) 5,790 -> 3,425, 6,343 -> 3,737, 11,465 -> 4,702, 14,281 -> 6,325. Plans of 1 to 64 tokens also favour it by about 10%. Those plans hold 3.3 to 6.8 pairs an item; a 4,096-token chunk holds about 9, so the bound is 12 (not timed at 4,096 directly; the served runs below use that step).
• Decode speed, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0: unchanged (4-bit 50.0 / 43.9 -> 50.2 / 44.1, oQ4e 42.1 / 36.2 -> 42.0 / 36.2 tok/s code / chat; decode takes the stream tiles).
• Prompt speed, tools/bench_cache.py --require-cold --tokens 16 on public 7,377- and 29,364-token arrays, ttft_s of the cold request after the warm-up: 4-bit 20.72 -> 19.90 s and 91.79 -> 88.48 s (-4.0% / -3.6%); oQ4e 22.42 -> 21.40 s and 95.37 -> 92.40 s (-4.5% / -3.1%). One run per cell. The routed products are a larger share of prompt time once the sparse attention is faster (HIP Flash Next: WMMA sparse attention for prompt chunks on gfx11 #1597), so the gain should grow with it.
• Tests run: zig build hip-kernel-contract on gfx1151: SUMMARY gfx1151 51 kernels: 51 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems (main: 50 kernels, the same 126 live-versus-recorded register differences for existing kernels under this machine's compiler); the new tile's cases (groups 32 and 64): non-finite, repeat, shapes and sums pass, byte-equal to the 128-row tile it replaces; live limits 176 VGPRs and 34,304 LDS bytes as recorded. zig build hip-kernel-test 106/106 (the routed prompt byte-equality test now compares the new tile against the matrix, GEMM and K-parallel tiles on every routed prompt cut). zig build test -Dhip=true --summary all 468/468 steps, 0 failed. python3 tools/zig/lean_check.py 426 on main 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 main, see HIP: dot2_tiles.hip fails to compile for gfx1100 with Ubuntu clang 21.1.8 ("branch size exceeds simm16") #1521); the records for gfx1030, gfx1100 and gfx1201 come from compiling the tile for each chip and reading its code object (gfx1030 and gfx1201 get a trapping stub, as they have no gfx11 matrix cores, and never select it). gfx1100 keeps its table unchanged (not measured there). 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.
• 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 @gilby. This meets the gate we posted. The item tile is in the routed byte-equality sweep against the matrix, GEMM and K-parallel tiles, served tokens are unchanged on both checkpoints, the pairs_per_item = 12 row keeps denser plans on today's tile, and cold prompts are 3% to 4.5% faster. 64 rows instead of the 16 or 32 we discussed is fine given your kernel times; a smaller tier can join row_tiles later. We'll land it with gfx1150, gfx1152 and gfx1153 listed as taking the table untested (same bytes). One note for next time: launches/affine.zig is now at exactly 600 lines, so the next change there will need to split it.

[gilby]

An update after #1663 landed and with #1651's 4,096-row chunks in view: the 64-row tile now runs only for prompt chunks of at most 1,024 rows (34356b1f2, rebased on 9e8c368b).

The pairs-per-item bound reads the plan's item capacity as the launch passes it (min(pairs, experts) + pairs / 8), which grows with the pairs, so a long prompt sits near 4 pairs an item at every chunk size and the bound of 12 admitted the tile everywhere. With #1663 the 128-row tile is faster once an expert holds a longer chunk's rows. Measured on gfx1151 with TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP (revision 2b170fa6), rocprofv3 kernel trace of a cold 7,869-token prompt, GPU window without / with the 64-row tile at each prefill_step (measured on main d377fcc6 with #1651 and the new sparse attention kernel applied, before this rebase):

prompt step	128-row tile only	with the 64-row tile
1,024	13.94 s	13.70 s
2,048	12.55 s	12.69 s
3,072	12.29 s	12.46 s
4,096	12.24 s	12.47 s

On the rebased head (main 9e8c368b, which has #1685 and the banded matrix tile), the 1,024-row case still gains a little: window 21.65 -> 21.44 s (-1.0%), the routed and dense GEMMs together 6.55 -> 6.34 s. The rebase also moves the item tile to GemmOut (#1685), so it writes what the 128-row tile writes; its contract runners stay bit-equal to that tile (15 shapes, worst 0.000) and SUMMARY gfx1151 54 kernels: 54 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 130 limit problems (main: 52 kernels, the same 2 and 130).

So the new fitsPrefillItem also asks for a chunk of at most 1,024 rows, and a host test checks gfx1151 picks the 64-row tile at 1,024 rows and the 128-row tile at 1,025 and 4,096 (it fails with the bound at 2,048). zig build test -Dhip=true --summary all 485/485 steps; lean 424 as on main. With #1651 the default chunk on gfx115x is 4,096 rows, so there this tile runs only under --policy prefill_step=1024 or less; if that makes it not worth carrying, I'm fine closing this instead.

[ashhart]

Thanks @gilby. Our review says this lands: 34356b1 caps the 64-row routed tile at prompt chunks of 1,024 rows, the tile is byte-equal to the 128-row one at every routed cut, and other GPUs are unchanged. But it conflicts with main now in zig/src/hip/launches/affine.zig. Could you rebase onto current main? We'll land it straight after.

[gilby]

Thanks @ashhart. Rebased onto 089a18df (a8ddcc6f3). The only conflict in affine.zig was two lines at the top: main added the stream_pick import where this change had dropped an unused Stream import and a blank line, so main's lines stay.

On the rebased head, gfx1151: zig build test -Dhip=true --summary all 507/507 steps, 0 failed; hip-kernel-test 115 passed, 1 skipped, 0 failed; python3 tools/zig/lean_check.py 424 as on main; SUMMARY gfx1151 54 kernels: 54 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 130 limit problems (main: 52 kernels, the same 2 and 130), and the item tile's records are current on all four chips.

With --policy prefill_step=1024, where the tile applies, TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6, --parallel 1 --context 131072 --prompt-cache-gib 0, one fresh server each (this change, then main): steady cold prompts 7,377 tokens 19.51 -> 19.24 s and 29,364 tokens 87.02 -> 85.85 s (-1.4% both), decode 51.05 / 44.99 -> 50.96 / 44.92 tok/s code / chat (level), token_sha256 identical to main across bench_openai, bench_concurrent and both bench_cache --require-cold prompts, MTP acceptance 2.753 / 2.37 in both. These are single runs, not alternated.

[ashhart]

Thanks @gilby: the push is a clean rebase of what we cleared (the 64-row item tile, gfx1151's table, the 1,024-row chunk cap, contract records and sweep), and other GPUs are unchanged. It's in the landing queue.

[gilby]

Rebased onto current main (cfcbecb), keeping your changelog commit as it was. The only conflict was the CHANGELOG, where main added entries beside it; both are kept. The three code commits applied unchanged.

Receipts on the rebased head (32dff8d), one Strix Halo iGPU (gfx1151), ROCm 7.1 hipcc, zig 0.17.0:

• zig fmt clean, lean_check 424 problems, the same as main.
• zig build test -Dhip=true --summary all: 516/516 steps, 6818 of 6867 tests passed, 49 skipped, 0 failed.
• hip-kernel-contract -Dhip-arch=gfx1151: SUMMARY gfx1151 59 kernels: 59 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed. The run exits 1 on 141 limit problems, which are stale records for kernels this change doesn't touch (main's were re-derived with ROCm 10.1, test(hip): re-derive every gate-5 record with ROCm 10.1 #1979; this box has 7.1).
• Device tests: 119 passed, 1 skipped, 0 failed, including the gfx1151 64-row selection test and the one that sends a routed prompt past its pairs per item to the 128-row tile.

Not run: gfx1030, gfx1100 and gfx1201 hardware, and any other family.