feat(quant): backend quantization capability table and load-time convert-or-reject policy
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801. Phase 1. Depends on #1802. #1807 and #1808 build on the table introduced here.
## Context
The ROCm backend handles quantization modes unevenly (measured on `gfx1151`, see #1801): affine 4/8-bit correct, mxfp8 correct after the dispatch fix in the mlxcelverse ROCm overlay, mxfp4 broken (qmm hang, GPU quantize launch failure at large shapes), NVFP4 unsupported (no group-16 dispatch, no E4M3 scale path).
mlxcel loads these modes through several paths that assume the backend can run them:
- NVFP4: `nvfp4_repack_strategy` (`src/models/sanitize.rs:533-547`) defaults to `DirectTranscode`, which produces native NVFP4 (group 16). The `DenseAffine` route already exists (re-quantizes to affine 4-bit, `sanitize.rs` around 850-868) but only when `MLXCEL_NVFP4_DENSE_REPACK=1` is set. Layers whose `in_dim` does not match the target group are skipped with a "Skipping NVFP4 repack" message (`sanitize.rs:662-709`) rather than failing.
- Vendor FP8 block checkpoints: `src/models/fp8_block.rs` always re-quantizes to mxfp8 (`MXFP8_MODE`, group 32, bits 8).
- mxfp4 checkpoints (for example gpt-oss): loaded as-is; no conversion path.
The only backend-capability policy today is the pre-Ampere CUDA one (`pre_ampere_cuda()` and `bf16_to_f16_at_load()`, `sanitize.rs:1803-1879`).
Direction for this epic: where the backend cannot run a mode, convert it to one it can run; reject at load with an actionable message only when conversion is impossible.
## Scope
A single place that answers "can this backend run quantization mode M natively?", and load-time policy that uses it to convert or reject. The conversions for mxfp8 and mxfp4 themselves live in #1807 and #1808; this issue covers the table, the policy hook, and NVFP4.
## Implementation plan
1. Add a backend quantization capability query (Rust side, fed by the backend kind from #1803 or the build features): per mode, `Native`, `ConvertTo(mode)` or `Unsupported`. Initial ROCm table: affine native; mxfp8 native (after #1807 validates it); mxfp4 `ConvertTo(affine 4-bit)` until #1808 fixes the kernels; NVFP4 `ConvertTo(affine 4-bit)`. Metal and CUDA report exactly today's behavior.
2. NVFP4: make `nvfp4_repack_strategy` pick `DenseAffine` by default when the backend reports NVFP4 as non-native. `MLXCEL_NVFP4_DENSE_REPACK` keeps working as an explicit override on every backend.
3. On a backend that needs conversion, a layer that cannot be converted (group mismatch, unexpected shape) fails the load with the layer name and reason instead of being skipped silently. Keep today's skip behavior on backends where the native path exists.
4. Log one line per load stating the source mode, the chosen route and the reason.
5. Measure conversion loss with the logit-trace method from `docs/benchmarks.md` ("Judging a change that moves the numbers"): the same NVFP4 checkpoint loaded natively on Metal versus through `DenseAffine` on Metal isolates conversion loss from backend differences; then `DenseAffine` on ROCm adds only the backend difference. Report decided-position mismatch rates in the PR body.
## Acceptance criteria
- [ ] An NVFP4 checkpoint loads and generates on ROCm with no environment variables, through `DenseAffine`, and the load log states the conversion.
- [ ] On ROCm, a layer that cannot be converted fails the load with a message naming the layer and the reason.
- [ ] Metal and CUDA choose the same routes as before for every mode (unit tests over the capability table and `nvfp4_repack_strategy`).
- [ ] The PR reports the conversion-loss logit-trace numbers described above.
## References
- `src/models/sanitize.rs` (`nvfp4_repack_strategy` 533-547, skip messages 662-709, `bf16_to_f16_at_load` 1803, `pre_ampere_cuda` 1879)
- `src/models/fp8_block.rs` (`MXFP8_MODE` line 68)
- Logit-trace method: `docs/benchmarks.md`, `examples/logit_trace`, `scripts/compare_logit_traces.py`
Contributor guide
Research direction
Start with src/models/sanitize.rs, especially nvfp4_repack_strategy, the NVFP4 skip path, bf16_to_f16_at_load, and pre_ampere_cuda; then inspect src/models/fp8_block.rs and the backend kind from #1803. Add capability-table and route-selection tests, and use docs/benchmarks.md, examples/logit_trace, and scripts/compare_logit_traces.py for validation. Done means ROCm converts or rejects with actionable logs while Metal and CUDA retain their existing routes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100