fix(rocm/quant): mxfp4 qmm hang and GPU quantize launch failure, with an affine 4-bit fallback
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801. Phase 2. Depends on #1802, #1804, #1806.
## 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);
- GPU `quantize` in mxfp4 is bit-exact to the CPU at 256x512 but fails at 4096x4096 with `hipLaunchKernel(...) failed: invalid configuration argument`;
- on the unfixed fork the mxfp4 path also triggered `HSA_STATUS_ERROR_MEMORY_FAULT` in `qmv_warp_shared_kernel<..., 4, 32, false, 16>`.
mlxcel has no conversion path for mxfp4 today: checkpoints are loaded as-is.
## 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 #1804 in place the failure should surface as an error rather than a hang, which makes bisecting practical.
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).
3. If native mxfp4 is fixed: flip it to `Native` in the #1806 table and add op checks for `quantized_matmul` and `gather_qmm`.
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).
5. Record any kernel fix in `patches-rocm/UPSTREAM` for upstreaming (#1813).
## 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.
- [ ] mxfp4 never hangs or faults the GPU on ROCm; unsupported cases fail with an error.
- [ ] If the fallback is used, the PR reports conversion loss (logit trace, Metal native versus Metal fallback) and load-time peak memory.
- [ ] Metal and CUDA mxfp4 behavior is unchanged.
## References
- `mlx/backend/rocm/quantized/qmm.hip` (qmv kernels and dispatch), `fp_quantize.hip`
- NVFP4 dense repack precedent: `src/models/sanitize.rs` (`Nvfp4RepackStrategy::DenseAffine`)
Contributor guide
Research direction
Start by reproducing the mxfp4 failures and read mlx/backend/rocm/quantized/qmm.hip, especially the qmv dispatch and fp4_e2m1_to_float path. Then inspect mlx/backend/rocm/quantized/fp_quantize.hip for the large-tensor launch failure and src/models/sanitize.rs for the NVFP4 DenseAffine precedent. Done means a supported mxfp4 checkpoint generates on ROCm without hangs or GPU faults, with unsupported cases failing clearly.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100