feat(rocm/quant): mxfp8 end-to-end on ROCm, including FP8 block checkpoints and the MoE gather path #1807

Part of #1801. Phase 2. Depends on #1802, #1806.

Status 2026-09-29: #1802 is closed (landed in PR #1818). #1806 is the current blocker: it owns the capability table that step 5 flips, and its NVFP4 abort is the only red item in the ROCm gate. That abort ends the mlxcel-core lib test binary at ffi_tests::compiled_qgelu_mlp_global_scale_native_nvfp4_prefill_matches_reference, so any mxfp8 op test added to mlxcel-core and sorting after that name will not run in make verify-rocm until #1806 is fixed; run it directly. The gather-path fix in step 1 has landed (PR #1818); what remains is committed tests, the end-to-end checkpoint run, and the table entry.

Context

mxfp8 matters for mlxcel beyond native mxfp8 checkpoints: every vendor FP8 block checkpoint is re-quantized to mxfp8 at load (src/models/fp8_block.rs, MXFP8_MODE), for example the Qwen3.5 FP8 text and VLM loaders.

The feasibility spike found mxfp8 quantized_matmul returning NaN on ROCm, and the unfixed fork faulting the GPU. The cause was in the ROCm qmv dispatch: kernels were instantiated with the activation dtype as ScaleT for every mode, while mxfp4/mxfp8 store one E8M0 byte per group, so the kernel read scales at 2-4x the real stride. mlxcelverse fixes this in mlx/backend/rocm/quantized/qmm.hip by dispatching non-affine modes with uint8_t scales. After the fix, on gfx1151:

• quantized_matmul mxfp8 matches an f32 dequantized reference for bf16 and f16 activations, M = 1, 4, 8, 64, up to 4096x4096 (relative error about 3e-3 for bf16, 3e-4 for f16);
• GPU quantize produces scales identical to the CPU and differs in 3.3% of weight bytes (tie rounding), with identical RMS error against the source weights (5.319e-4);
• dequantize on the GPU is bit-exact to the CPU.

Not yet verified: gather_qmm (the MoE path) in fp modes, and a real FP8 block checkpoint end to end.

Update 2026-09-29 (overlay changes since filing): the overlay landed in this repo with PR #1818 as src/lib/mlx-cpp/patches-rocm/, carrying the scale-type fix above as LOCAL_FIXES.md item 8. It also added two mxfp-relevant fixes found while validating: item 10, GatherQMM::eval_gpu sent non-affine modes that missed the affine fast paths to a gather_qmv_kernel instantiated with AFFINE=true and T-typed scales, so fp weights were dequantized with the affine formula (relative error 1.0 to 1e34 for mxfp8); non-affine modes now launch gather_qmv_kernel<T, uint8_t, BITS, 32, false> and other fp group sizes throw. And item 9, the expert-batched gather qmv path returned wrong bf16 results with sorted indices and is now opt-in via MLX_ROCM_GATHER_QMV_EXPERT_BATCHED=1. PR #1818 reports mxfp8 and mxfp4 gather_qmm matching a dequantized f32 reference (about 3e-3 relative for bf16) with sorted and unsorted indices, from an ad-hoc probe rather than a committed test. Later overlay commits (PRs #1822, #1823, #1824, #1856, #1861) touch affine-only qmv paths, scatter, indexing, Hadamard, SearchSorted and FFT; none changes mxfp8 handling. docs/installation.md:476 now tells users mxfp8 runs natively on ROCm.

Update 2026-09-29 (which MoE path runs on ROCm): the fused MoE launchers are in the custom-kernel port table with no ROCm port: moe_gateup_ports() and moe_down_ports() (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:2163, :2201) and moe_fc1_relu2_ports() (:2278) all set .rocm = nullptr, and select_kernel_port refuses them with a catchable error instead of aborting (PRs #2018, #2026, #2029). fused_moe_enabled() (src/models/switch_layers.rs:88-104) also requires custom_kernels_available(), which is false on ROCm (src/lib/mlx-cpp/turbo/gpu_backend.h:50-53). So the MLX graph path through gather_qmm is what an mxfp8 MoE runs on ROCm today, and it is the path to test. MLXCEL_FUSED_MOE=0 is not needed.

Scope

Prove and, where needed, fix mxfp8 on ROCm for the paths mlxcel actually uses, then mark it Native in the capability table from #1806.

Implementation plan

1. Add op-level checks for gather_qmm in mxfp8 (MoE expert gather, sorted and unsorted indices) against a dequantized reference. Apply the same uint8_t scale dispatch fix to the gather path if it has the same bug. (2026-09-29: the fix is done, LOCAL_FIXES.md item 10 in PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818; the gather path had a related but different bug, an AFFINE=true instantiation. The committed op checks are still missing: a search of src/, tests/ and examples/ found no test that checks gather_qmm numerically in mxfp8 (the mxfp8 cases in switch_layers.rs and the model tests only validate loader rejection). Cover sorted and unsorted indices with bf16 and f16 activations, and state whether the check exercises the default path or the opt-in expert-batched one from item 9.)
2. Load and generate with a vendor FP8 block checkpoint through fp8_block.rs on ROCm. Compare logit traces with the same checkpoint on Metal (test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809 provides the Metal baseline). (2026-09-29: not done. The gfx1151 correctness matrix, PR test(rocm): add the gfx1151 correctness matrix and its gate #1826, docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md, traced affine 4-bit only; its line 174 assigns mxfp8 coverage to this issue.)
3. Decide whether GPU-side quantization at load needs to be bit-identical to the CPU. If the tie-rounding difference moves decided positions, quantize on the CPU stream at load; otherwise document the difference.
4. Add a round-trip test for the ROCm path next to the existing fp8_block_requantize_round_trip_stays_within_half_an_e4m3_step (src/models/fp8_block_tests.rs:261). (2026-09-29: partly covered. That test and its neighbours in fp8_block_tests.rs are not backend-gated, and they sit in the mlxcel lib target, which passes in make verify-rocm on d8d34e2b, so the existing round trip already runs on the ROCm default device. A separate ROCm test is needed only if step 3 moves load-time quantization to the CPU stream and that choice needs pinning.)
5. Flip mxfp8 to Native for ROCm in the capability table. (2026-09-29: blocked on feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806, which has not created the table yet.)

Acceptance criteria

• gather_qmm mxfp8 op checks pass on ROCm. (PR test(rocm): prove mxfp8 end to end, incl. FP8 block checkpoints #2071: models::switch_layers::mxfp_tests.) (2026-09-29: the underlying fix and an ad-hoc check are in PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818; the committed test is missing.)
• An FP8 block checkpoint generates on ROCm with a decided-position mismatch rate against Metal within the threshold used in test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809. (PR test(rocm): prove mxfp8 end to end, incl. FP8 block checkpoints #2071: against the CPU device, Metal being unavailable on the ROCm host; 0 of 50 decided positions disagree. 2026-09-29: that threshold is --decided 2.0, documented in docs/benchmark_results/rocm-correctness-gfx1151-2026-09-12.md, PR test(rocm): add the gfx1151 correctness matrix and its gate #1826.)
• The ROCm round-trip test passes; the existing Metal/CUDA tests are unchanged.

References

• src/models/fp8_block.rs, src/models/fp8_block_tests.rs:261
• ROCm qmv dispatch: mlx/backend/rocm/quantized/qmm.hip (DISPATCH_GROUP_SIZE, launch_qmv); in this repo src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip
• MXFP8 scale rounding that mlxcel depends on: Round mxfp8 block scales up to avoid saturation ml-explore/mlx#4353
• (Added 2026-09-29) Overlay fixes: src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md items 8, 9 and 10; fused MoE port tables in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp; port dispatch reasoning in TECHNICAL_REPORTS/2026-kernel-port-dispatch-standard-20260929.*

Refresh log

2026-09-29

• Recorded feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 as the current blocker: it owns the capability table, and its NVFP4 abort hides later-sorting mlxcel-core tests from the ROCm gate.
• Marked the gather_qmm fix in step 1 as landed (PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818, LOCAL_FIXES.md item 10) and noted item 9's opt-in expert-batched path; the committed op test remains open.
• Recorded that the fused MoE launchers have .rocm = nullptr and refuse with a catchable error through select_kernel_port (PRs fix(rocm): guard the four remaining custom-kernel launchers #2018, refactor(core): choose custom-kernel ports through one helper #2026, refactor: finish the kernel port standardization #2029), so the gather_qmm graph path is what runs on ROCm.
• Checked every overlay commit since filing: none after PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818 changes mxfp8 handling. Marked step 4 partly covered by the existing, ungated round-trip test that already runs in the passing mlxcel lib target, and cited test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809's --decided 2.0 threshold.