fix(rocm/quant): mxfp4 qmm hang and GPU quantize launch failure, with an affine 4-bit fallback #1808

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

Status 2026-09-29: most of the kernel work this issue asked for landed with PR #1818, the PR that vendored the ROCm overlay (merged a few hours after this issue was filed). That PR reports mxfp4 quantized_matmul and gather_qmm matching a dequantized f32 reference, mxfp4 GPU quantize bit-identical to the CPU at 4096x4096, and gpt-oss-20b-MXFP4-Q4 generating coherent text natively. What remains: confirm that the original hang does not reproduce, commit op tests, and flip the capability table. #1802 is closed (PR #1818). #1806 is the current blocker: it owns the table and the load-log policy the acceptance criteria need, and its NVFP4 abort is the only red item in the ROCm gate. #1804 is open and not started; with no mxfp4 hang left to bisect, it no longer blocks the work here.

Context

mxfp4 is the native format of checkpoints such as gpt-oss (which ship mxfp4 MoE experts). On gfx1151 with the mlxcelverse ROCm overlay (including the E8M0 scale-type dispatch fix that repaired mxfp8):

• quantized_matmul in mxfp4 hangs even at 256x512 with M = 1 (process killed by timeout); (2026-09-29: not reproduced in PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818's validation, which reports mxfp4 quantized_matmul matching a dequantized f32 reference at about 3e-3 relative for bf16. No LOCAL_FIXES.md item names a hang fix, so which change removed it is not recorded; rerun this exact case before closing.)
• GPU quantize in mxfp4 is bit-exact to the CPU at 256x512 but fails at 4096x4096 with hipLaunchKernel(...) failed: invalid configuration argument; (fixed in PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818: the failing launch was arg_reduce, not fp_quantize.hip. For short reduction axes it launched one 1024-thread block per output on a 1-D grid, and AMD limits a grid dimension to 2^32 - 1 threads, which a 4096x4096 argmin over 16.7M elements exceeds. patches-rocm/LOCAL_FIXES.md item 11.)
• on the unfixed fork the mxfp4 path also triggered HSA_STATUS_ERROR_MEMORY_FAULT in qmv_warp_shared_kernel<..., 4, 32, false, 16>. (Fixed by the E8M0 scale-type dispatch, LOCAL_FIXES.md item 8, PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818.)

mlxcel has no conversion path for mxfp4 today: checkpoints are loaded as-is.

Update 2026-09-29 (overlay changes for mxfp4 since filing): besides items 8 and 11, PR #1818 fixed gather_qmm for mxfp4 and mxfp8 (LOCAL_FIXES.md item 10): non-affine modes that missed the affine fast paths reached a gather_qmv_kernel instantiated with AFFINE=true and T-typed scales; they now launch gather_qmv_kernel<T, uint8_t, BITS, 32, false>, and other fp group sizes throw. It also made the expert-batched gather qmv path opt-in (MLX_ROCM_GATHER_QMV_EXPERT_BATCHED=1, item 9) after it returned wrong bf16 results with sorted indices. Later overlay commits (PRs #1822, #1823, #1824, #1856, #1861) do not change mxfp4 handling. docs/installation.md:476 now tells users mxfp4 runs natively on ROCm, including gpt-oss-20b-MXFP4-Q4.

Update 2026-09-29 (which MoE path runs): the fused MoE launchers have no ROCm port (.rocm = nullptr in moe_gateup_ports(), moe_down_ports() and moe_fc1_relu2_ports(), src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:2163, :2201, :2278) and refuse with a catchable error through select_kernel_port (PRs #2018, #2026, #2029). fused_moe_enabled() (src/models/switch_layers.rs:88-104) is false on ROCm because custom_kernels_available() is. mxfp4 MoE experts therefore run through the gather_qmm graph path. PR #1818 measured gpt-oss-20b decode at 3.6 tok/s through the generic gather kernel; speed is #1814's concern, not this issue's.

Update 2026-09-29 (committed evidence and its limit): ffi_tests has four mxfp4 group-32 cases that run by default (plus the #[ignore]d shapeless_compile_audit_harness). Two sort before the NVFP4 abort tracked in #1806 and run in the ROCm gate, whose only failure in -p mlxcel-core --lib is that abort: compiled_qgelu_mlp_global_scale_all_present_single_token_matches_reference (src/lib/mlxcel-core/src/ffi_tests.rs:2836, M = 1) and compiled_qgelu_mlp_global_scale_mixed_down_only_matches_reference (:3045). Their shapes are 128 and 256 wide, below the 256x512 hang case. The other two, ..._non_nvfp4_multi_token_fallback_matches_reference (:2916) and ..._no_sidecar_matches_unscaled (:3105), sort after the abort and never run until #1806 is fixed; run them directly.

Scope

Make mxfp4 checkpoints usable on ROCm. Preferred: fix the kernels so mxfp4 runs natively. Fallback, per the epic's direction: convert mxfp4 to affine 4-bit at load when the backend reports mxfp4 as non-native.

Implementation plan

1. Root-cause the qmm hang. Candidates to check first: the 4-bit fp unpack path in the qmv kernels (fp4_e2m1_to_float), the tiled or idot dispatch conditions that assume affine, and qmv_t/gather_qmv variants reached for bits = 4 in fp mode. With fix(rocm): surface ROCm GPU failures through Event::error instead of NaN or hangs #1804 in place the failure should surface as an error rather than a hang, which makes bisecting practical. (2026-09-29: first reproduce it. Run mxfp4 quantized_matmul at 256x512, M = 1, with a timeout, on the current overlay. If it passes, record the result here and drop the root-cause work; if it hangs, this step stands as written.)
2. Root-cause the quantize launch failure in mlx/backend/rocm/quantized/fp_quantize.hip (grid or block computation for large tensors; check the path that actually serves mode="mxfp4" since the scales come out correct at small shapes). Done in PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818: the cause was arg_reduce launch geometry (LOCAL_FIXES.md item 11), not fp_quantize.hip.
3. If native mxfp4 is fixed: flip it to Native in the feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 table and add op checks for quantized_matmul and gather_qmm. (2026-09-29: this is now the expected path. The op checks were run ad hoc in PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818 but not committed; no committed test checks mxfp4 gather_qmm numerically. The table flip waits on feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806.)
4. If not, or as an interim: add an mxfp4 to affine 4-bit dense repack at load, modeled on the NVFP4 DenseAffine path (src/models/sanitize.rs). Choose the affine group from in_dim (32/64/128), fail the load for layers that cannot be converted, and measure load-time peak memory on a large MoE (the dense intermediate is the cost). (2026-09-29: likely obsolete. Needed only if step 1 reproduces the hang; kept for that case.)
5. Record any kernel fix in patches-rocm/UPSTREAM for upstreaming (chore(rocm): mlxcelverse ROCm fork sync script, MLX pin-bump procedure, and upstreaming local fixes #1813). (Corrected 2026-09-29: local fixes are recorded in src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md; UPSTREAM only pins the fork commit. Items 8, 10 and 11 are already recorded there.)

Acceptance criteria

• An mxfp4 checkpoint (for example gpt-oss-20b) generates on ROCm, natively or through the affine fallback, and the load log states which. (2026-09-29: generation is evidenced natively by PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818; the load-log line depends on feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806's policy hook.)
• mxfp4 never hangs or faults the GPU on ROCm; unsupported cases fail with an error. (2026-09-29: the known hang and launch failure are gone per PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818; fp group sizes other than 32 in gather_qmm now throw instead of computing garbage. The hang case is not yet re-run on record.)
• If the fallback is used, the PR reports conversion loss (logit trace, Metal native versus Metal fallback) and load-time peak memory. (2026-09-29: moot unless step 4 is needed.)
• Metal and CUDA mxfp4 behavior is unchanged.

References

• mlx/backend/rocm/quantized/qmm.hip (qmv kernels and dispatch), fp_quantize.hip; in this repo under src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/
• NVFP4 dense repack precedent: src/models/sanitize.rs (Nvfp4RepackStrategy::DenseAffine)
• (Added 2026-09-29) src/lib/mlx-cpp/patches-rocm/LOCAL_FIXES.md items 8 to 11; PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818 validation notes

Refresh log

2026-09-29

• Recorded that PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818 landed the mxfp4 kernel fixes (LOCAL_FIXES.md items 8, 10, 11) and reports mxfp4 qmm, gather_qmm and quantize correct and gpt-oss-20b generating natively; checked off step 2 with its real root cause (arg_reduce grid geometry).
• Recorded feat(quant): backend quantization capability table and load-time convert-or-reject policy #1806 as the current blocker (capability table, load-log policy, and the gate abort that hides two of the four mxfp4 ffi_tests cases); noted that fix(rocm): surface ROCm GPU failures through Event::error instead of NaN or hangs #1804 no longer blocks this work.
• Turned step 1 into a reproduce-first check, since no record names the change that removed the hang, and marked the affine fallback in step 4 as likely obsolete rather than deleting it.
• Recorded that fused MoE refuses on ROCm with a catchable error, so mxfp4 experts run through the gather_qmm graph path, and that no overlay commit after PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818 touches mxfp4.
• Corrected the record file for kernel fixes (LOCAL_FIXES.md, not UPSTREAM).