lablup / lablup/mlxcel

feat(quant): backend quantization capability table and load-time convert-or-reject policy

Open
#1,806 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 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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.