lablup / lablup/mlxcel

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

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.