lablup / lablup/mlxcel

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

Open
#1,808 0 comments 0 reactions 0 assignees View on GitHub
arch:moe area:models platform:linux priority:medium status:backlog type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.