Add ROCm MXFP4 routed expert kernels

[jyatesdotdev]

Summary

• add resident ROCm MXFP4 x Q8_K kernels for routed expert gate/up and down projections
• coalesce decode weight reads across wave32 lanes and unpack MXFP4 nibbles with register permutations
• keep gate/up and down reductions register-resident on gfx1151
• route two-to-four-token MXFP4 batches through the direct kernels instead of partially filled expert tiles
• extend the independent CPU-oracle ROCm test across 1, 2, 3, 4, 5, 32, 128, and 512 tokens
• report MXFP4 routed weights as 4-bit metadata and compile Strix Halo host objects with PIC

This is intentionally stacked on ds4f-mxfp4. It does not include SSD-selected-expert changes, MTP or distributed protocol work, or the Q8 prequant restoration already covered by #623.

Correctness

The focused head 7d63867 was built and tested independently on both Strix Halo systems, max and max2, with Radeon 8060S Graphics / gfx1151:

make -j1 strix-halo HIPCC=/opt/rocm-therock/bin/hipcc
LD_LIBRARY_PATH=/opt/rocm-therock/lib make test-mxfp4-rocm HIPCC=/opt/rocm-therock/bin/hipcc

Both full ROCm builds passed. Every synthetic gate/up, SwiGLU, and down result matched the CPU oracle with zero tolerance failures at all eight token counts. Worst observed absolute error was 2.15e-6. Counts 2 through 4 cover the new direct tiny-batch route, while count 5 verifies the transition back to expert-sorted tiles.

The optimization was also checked against the production Flash model. Baseline and optimized first-frontier logits were byte-identical across 129,280 values, including the same argmax.

Default-backend checks passed:

• macOS build of ds4, ds4-server, ds4-bench, ds4-eval, and ds4-agent
• ds4-eval extractor self-tests and ds4-agent tests
• layer packing: 97/97
• multi-GPU placement: 98/98
• GPU argument and CLI tests
• scalar MXFP4 dot test: 4/4

Full-model performance

These are controlled A/B measurements from the production integration tree with the same 0731 MXFP4 model, dual-host 0:21 / 22:output split, resident weights, and runtime configuration. Only the MXFP4 kernel variant or tiny-batch dispatch changed between each before and after point.

Workload	Before	Optimized	Change
Single-token decode	11.35 tok/s	14.56 tok/s	+28.3%
Two-token aggregate	11.13 tok/s	14.88 tok/s	+33.7%
Four-token aggregate	17.01 tok/s	23.61 tok/s	+38.8%

A live OpenAI-compatible request after deployment decoded at 14.86 tok/s. The single-token gate kernel uses 52 VGPR / 38 SGPR and the down kernel uses 32 VGPR / 24 SGPR; neither spills or uses private memory. The batched down specialization uses 32 VGPR / 27 SGPR with no spills.

The performance run uses the Q8 prequant path from #623; disabling that path remains substantially slower. This PR does not duplicate #623.

Q4_K comparison on the same PR head

To keep the causal scopes separate, the MXFP4 parent-versus-head A/B above measures this PR. The following compares the Q4_K and MXFP4 routed-expert format paths using the same exact build from head 7d63867.

A tensor-level audit of the two 0731 artifacts found identical tensor names and shapes. All 1,199 non-routed tensor payloads were byte-identical; the remaining 129 tensors were the routed-expert weights encoded as Q4_K or MXFP4. All 58 common metadata records were also byte-identical; Q4_K carries four additional imatrix-provenance records.

Metric	Q4_K	MXFP4	MXFP4 vs Q4_K
GGUF size	153.33 GiB	145.26 GiB	-5.26%
Planner model allocation, coordinator / worker	79.16 / 74.69 GiB	75.03 / 70.75 GiB	-4.13 / -3.94 GiB
Frozen strict sanity checks	23/24	23/24	same pass set
4K prefill	72.11 tok/s	75.19 tok/s	+4.3%
4K steady decode	7.80 tok/s	8.17 tok/s	+4.7%
First generated token	128.403 ms	122.529 ms	-4.6%

Performance values are medians of three alternating runs per format with a fixed 4K prefill and 256-token decode, warm resident weights loaded from local SSD, fresh processes and format-specific KV state, the same dual-host 0:21 / 22:output split, and MTP disabled. Startup time is excluded.

A single paired 32K control reproduced the same direction and magnitude: MXFP4 improved prefill by 4.3%, steady decode by 4.7%, and first-token latency by 4.5%. Because this was one pair rather than three alternating pairs, it is treated only as a corroborating long-context check.

These absolute performance values are not directly comparable to the production-integration figures above: this section uses the exact PR-head build and a different fixed-token benchmark shape. It compares format paths and does not attribute the Q4_K-to-MXFP4 difference to this PR.

The 24-case suite is a targeted regression check, not evidence that either format is generally higher quality or that the formats are equivalent. Both passed the same 23 cases, and neither produced a repetition failure.

[jyatesdotdev]

I have two strix halos and wanted to run the MXFP4 variant, so spent a decent amount of time working with Codex (Sol 5.6) to get this working. I have more unrelated changes for getting MTP working across the systems and I'd like to see how difficult it would be to enable RDMA over thunderbolt (though that later part involves way more than just this repository). Let me know if you'd like any changes for this PR!

[jyatesdotdev]

Also note the current numbers are from iommu=pt, I'm testing now with it off, which should produce better numbers. I'm going to be daily driving this on my own hardware for a while, so I'll report back if there are any quality issues. I'm using pi for my harness.

[jyatesdotdev]

Follow-up to my earlier IOMMU note: amd_iommu=off produced a modest, repeatable speedup.

These are optimized deployment results, not standalone #656 throughput. The tested build combines the resident ROCm MXFP4 routed-expert kernels from this PR with the restored Q8 prequant decode path from #623. No #623 code is included in this PR.

The A/B held the binaries, local-NVMe GGUF, Fedora kernel 7.1.5-101.fc43, model split, and benchmark flags fixed across the reboot. The only configuration change was iommu=pt versus amd_iommu=off; both retained the same 126,976 MiB GTT and page settings.

• two Radeon 8060S / gfx1151 systems
• layers 0:21 / 22:output
• 32-bit distributed activations
• 4096-token local/distributed prefill chunks, window 3
• 128 greedy decode tokens at 2K, 8K, and 32K
• one warm-up followed by three measured runs per configuration
• table reports medians

Context	iommu=pt prefill / decode	amd_iommu=off prefill / decode	Prefill / decode delta
2K	67.70 / 10.67 tok/s	69.07 / 10.95 tok/s	+2.0% / +2.6%
8K	84.08 / 10.15 tok/s	88.40 / 10.47 tok/s	+5.1% / +3.2%
32K	114.88 / 9.58 tok/s	120.99 / 9.93 tok/s	+5.3% / +3.7%

All 18 dumped frontier-logit files were byte-identical at each matching context across every repetition and both boot modes, so this affected performance rather than model output.

On these systems, amd_iommu=off disables the NPU/SVA path. The result points to a roughly 3% platform-level decode gain and up to 5% longer-prefill gain; it does not indicate a change in the MXFP4 kernel itself.

[jyatesdotdev]

Follow-up on commit 7d63867, which contains the focused MXFP4 optimization:

• coalesces MXFP4 weight reads across wave32 lanes and performs nibble unpacking with register permutations;
• keeps gate/up, SwiGLU, and down reductions register-resident on gfx1151;
• routes 2–4 token batches through direct kernels instead of partially filled expert-sorted tiles;
• introduces no register spills or private-memory use.

The CPU-oracle test now covers 1, 2, 3, 4, 5, 32, 128, and 512 tokens. All cases passed, with worst observed absolute error of 2.15e-6; production first-frontier logits were byte-identical across 129,280 values.

Controlled parent/head results were:

• single-token decode: 11.35 → 14.56 tok/s (+28.3%);
• two-token aggregate: 11.13 → 14.88 tok/s (+33.7%);
• four-token aggregate: 17.01 → 23.61 tok/s (+38.8%).

The Q4_K and IOMMU measurements in the PR are additional format/platform context and are not attributed to this commit.

[jyatesdotdev]

@antirez any chance for a review or comment? Curious to know what you think and if I should spend some more time improving this! Thanks!