feat(rocm/quant): mxfp8 end-to-end on ROCm, including FP8 block checkpoints and the MoE gather path
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801. Phase 2. Depends on #1802, #1806.
## 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.
## 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.
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 (#1809 provides the Metal baseline).
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`).
5. Flip mxfp8 to `Native` for ROCm in the capability table.
## Acceptance criteria
- [ ] `gather_qmm` mxfp8 op checks pass on ROCm.
- [ ] An FP8 block checkpoint generates on ROCm with a decided-position mismatch rate against Metal within the threshold used in #1809.
- [ ] 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`)
- MXFP8 scale rounding that mlxcel depends on: ml-explore/mlx#4353
Contributor guide
Research direction
Start with the ROCm dispatch in mlx/backend/rocm/quantized/qmm.hip and the existing fp8_block round-trip test at src/models/fp8_block_tests.rs:261; run the relevant quantized_matmul and gather_qmm checks. Then trace the checkpoint path through src/models/fp8_block.rs and compare generation with the Metal baseline from #1809. Done means ROCm gather_qmm checks, checkpoint generation, and the ROCm round-trip test meet the stated acceptance criteria and the capability table marks mxfp8 Native.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100