Repository navigation
HIP: bound the routed prompt's K-parallel tile by pairs per item, and give it gfx1201's short prompts - #1607
Merged
Merged
Conversation
A routed prompt's K-parallel tile on gfx1201 wins or loses by the pairs a plan holds per item, which its items alone cannot tell. A tuning row can now say the most pairs per item it takes, and the registry refuses a plan with more, given the plan's pairs in its shape. No table sets the field yet, so every choice stays the same; a test keeps the field to routed prompt rows.
The registry can now bound a routed prompt by its pairs per item, so serving passes each plan's pairs to the prompt tile's choice beside its items, host-side only: the kernels' arguments keep their layout. The affine fixtures pass theirs too, pickable sweeps a routed prompt plan's pair counts, and the prompt path's byte check now compares every launched entry, since all of them are the gemm family.
The tuner's routed cells held only Qwen3.6's geometry, so it could not see that Nemotron's plans run slower on the K-parallel tile. Its classes now carry each model's plan: Qwen3.6, Nemotron and Flash Next with their experts, slots, shared slot and bucketed or exact sizing, and a served plan without a shared slot draws distinct experts a token, as serving does. Existing cases draw and digest as before.
Moving a routed prompt's items bound alone cannot both give Qwen3.6 the K-parallel tile and keep Nemotron's full plans off it, and when kp3 refuses a plan, kp1 takes it first. The proposal now lives in a host-tested search that moves kp3's items and pairs bound and kp1's pairs bound together, scored with the per-cell veto on each whole table, and the tuner times both sides of every pairs bound it tries.
Served runs on the R9700 put the routed K-parallel tile's threshold between 3.43 and 3.75 pairs per item, which no whole number holds, so the bound is now an f32 compared exactly in f64. The tuner tries the least bound that takes each timed cell, times both sides of every place its proposal changes a choice, and retakes a noisy timing, so one noisy cell no longer vetoes a whole bucket by chance.
On gfx1201 the routed prompt's K-parallel tile beats the GEMM tile for plans that hold few pairs an item and loses for full ones. Its row now takes plans of up to 507 items and 3.5 pairs an item, from served runs on Qwen3.6 and Nemotron: Qwen3.6's 33 to 128-token prompts fill x0.53 to x0.80 and Nemotron's 129 to 186-token ones x0.82 to x1.00, other sizes stay within noise, and every reply keeps its bytes.
With gfx1201 now shipping a pairs bound, the joint search's tests and the entries' unbounded control started from a table that already set it, so they could pass without the search or the field doing their work. They now start from gfx1100's rows, the control probes above 3.5 pairs an item, and the registry's and gfx1201's comments say what the final bound does and where its 3.5 came from.
ashhart
force-pushed
the
hip-routed-pairs
branch
from
October 29, 2026 23:45
4399e2d to
37680bc
Compare
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.
Fixes #1557.
What this changes
On gfx1201 the routed prompt's K-parallel tile (
mlx.routed.prefill.kp3) beats the dot2 GEMM tile for plans that hold few pairs an item and loses for full ones, but the registry could bound a routed plan only by itsitems, which is a bucket's capacity, so #1527 had to keep gfx1100's 300. This PR:tuning.Row.pairs_per_item, anf32, 0 off;registry.Shape.plan_pairs):Registry.costrefuses a routed plan of more thanpairs_per_item * itemspairs, compared inf64, exactly for any item count below 2^29. A tuning test keeps the field tomlx.routed.prefill.*rows, finite and from 0 to 8. No table but gfx1201's sets it, so no other GPU's choice changes.prefillLaunchbeside its items, host-side only:affine.ArgandRoutingkeep their layout (their size assert stands). The affine fixtures pass theirs,pickablesweeps a routed prompt plan's pair counts, and the prompt path's byte check now compares every launched entry (all of them are the gemm family), so kp3 and the GEMM tile are compared on each routed prompt plan.hip-affine-tunetimes every routed MLX model's plans: Qwen3.6 (256 + shared, 9 slots), Nemotron (128, 6) and Flash Next (512 + shared, 11, exact sizing), and a served plan without a shared slot draws distinct experts a token. Its proposal moved intozig/src/hip/tuning/search.zig, host-tested: kp3's items and pairs bound and kp1's pairs bound are searched together (refusing kp3 hands a plan to kp1 before the GEMM tile), candidates are each cell's least admittingf32, the tuner times both sides of every place its proposal changes a choice, and a noisy timing is retaken up to twice.items = 507,pairs_per_item = 3.5. The tuner proposed 507 and 4.14; served runs lowered the bound to 3.5 (below).No kernel or arithmetic change: every routed prompt entry is in the gemm family, and replies keep their bytes in every run below.
Why 3.5, and why a fraction
The tuner times each tile in isolation on plans whose experts are drawn uniformly; a trained router concentrates pairs on popular experts, so served items hold more rows, and the tuner admits kp3 further than serving can take it (at Nemotron's 213 tokens it timed kp3 faster, where served it is 8% slower). Served runs on this branch's code, alternated with
mainon fresh boots:mainmainmain: x1.003 to x1.005main: x1.000 / x1.001No whole number lies between 3.43 and 3.75, so the bound is a fraction; 3.5 keeps Nemotron's 65 to 128-token plans on kp3 as today, gives it kp3 from 129 to 186 tokens, and keeps 187 to 256 on the GEMM tile.
Receipt
upstream/mainat1017c5c6, Zig 0.17.0, HIP 7.16 (AMD clang 24.0.0git), one Radeon AI PRO R9700 (gfx1201) that nothing else was using, Linux. Checkpointsmlx-community/Qwen3.6-35B-A3B-4bitat38740b847e4cb78f352aba30aa41c76e08e6eb46andTensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bitatd9d758fb83953437f7263256b0d96157e2a348b8.zig build test -Dhip=true466/466 steps, 5687/5716 tests passed (29 skipped);hip-kernel-test106/106;hip-affine-test30/30;tf-qwen35-test checkon Qwen3.6: 42 passed, 0 failed, 5 skipped.tools/zig/lean_check.py: 425 problems onupstream/mainand on this branch, none new.atc overlap: no conflict.upstream/mainat7c068895with it plus this PR's code commits and row;7c068895..1017c5c6touches no HIP, lane or registry file. Server with--prompt-cache-gib 0, a fresh process each round, rounds alternating main, branch, main, branch, then a reversed pair. Cold prefill:tools/bench_cache.py --require-cold --tokens 16 --repeats 3 --streams 1, medianttft_sa round. Every reply'stoken_sha256(cold prompts) andtoken_sha(greedy and seededtools/bench_openai.pyreplies, temperatures 0 and 1.0) is equal between the builds in every round.main(tables below).zig build hip-affine-tuneon the R9700 (raw routed times below):BOUND mlx.routed.prefill kp3 items 300 -> 507, kp3 pairs_per_item 0 -> 4.143750190734863, kp1 pairs_per_item 0 -> 0: 56 cells changed, loss 0.000, saving 91796.99 us, search 5333 ms; every plain prompt bound keeps gfx1100's value.Not run
Flash Next (
TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP, 73.3 GiB) needs three of these 32 GB cards, and the others here serve other work. The tuner covers its plan geometry. Its exact plans hold 0.89 pairs an item up to 513 pairs, where Qwen3.6 gains most, anditems = 507gives it kp3 up to 41 rows (24 today). If you'd rather it waited for a served Flash Next run, the last two commits (perf(hip): gfx1201's routed prompt tiles by pairs per itemand the test fix after it) drop together, and the rest stands on its own: no other table sets the field. gfx1200 takes the table unmeasured, as before.Qwen3.6-35B-A3B, cold prefill (median
ttft_sof three cold requests a boot)Nemotron 3.5 Lightning, cold prefill
Qwen3.6, the size after a moved bound, in boots of its own
Nemotron, the sizes after a moved bound, in boots of their own
Nemotron decode, ten alternated boots a build (
tools/bench_openai.py --tokens 128 --reps 3, tok/s)The tuner's routed prompt times from its last run: median us a launch (noise when above 5%); U marks a cell whose table copies could not reach 512 MiB a rotation
35b served gate_up (1024 x 2048, group 64)
35b served down (2048 x 512, group 64)
nemotron served up (1856 x 2688, group 64)
nemotron served down (2688 x 1856, group 64)
flash next served gate_up (1280 x 2560, group 32)
flash next served down (2560 x 640, group 32)
35b served gate_up (1024 x 2048, group 64)
35b served down (2048 x 512, group 64)
nemotron served up (1856 x 2688, group 64)
nemotron served down (2688 x 1856, group 64)
flash next served gate_up (1280 x 2560, group 32)
flash next served down (2560 x 640, group 32)
nemotron served up (1856 x 2688, group 64)
nemotron served down (2688 x 1856, group 64)
flash next served gate_up (1280 x 2560, group 32)
flash next served down (2560 x 640, group 32)
nemotron served up (1856 x 2688, group 64)
nemotron served down (2688 x 1856, group 64)
flash next served gate_up (1280 x 2560, group 32)
flash next served down (2560 x 640, group 32)
flash next served gate_up (1280 x 2560, group 32)
flash next served down (2560 x 640, group 32)