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

Part of #1801. Phase 1. Depends on #1802 (closed; landed in PR #1818). #1807 and #1808 build on the table introduced here.

Status 2026-09-29: this is the last red item in the ROCm gate and the effective blocker for #1807, #1808 and #1809. make verify-rocm on the gfx1151 host at d8d34e2b (smoke step excluded) has exactly one failing target, -p mlxcel-core --lib, and exactly one terminate called, both from ffi_tests::compiled_qgelu_mlp_global_scale_native_nvfp4_prefill_matches_reference (src/lib/mlxcel-core/src/ffi_tests.rs:2980). Every other step passes (crate versions, kernel dtype keys, kernel port dispatch, llama-compat manifest, cargo fmt --check, workspace clippy with --features rocm). Implementation has not started.

Gate impact (added 2026-09-29)

The test runs native NVFP4 (group 16, bits 4) through compiled_gelu_approx_mlp_forward_global_scale and evaluates with eval; neither returns Result across the cxx bridge, so an exception thrown by the ROCm backend terminates the process. The ROCm qmv dispatch only has group sizes 32, 64 and 128 (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip:4277-4293) and throws for anything else. The exact message printed by the gate run was not captured in this refresh.

verify-test-rocm runs single-threaded (--test-threads=1, Makefile:782) and libtest runs tests in name order, so the abort ends the mlxcel-core lib test binary: no test sorting after that name reaches the runner at all. "Not in the failure list" does not mean "passed". Tests never reached include the rest of ffi_tests after that name (among them the second native-NVFP4 test, quantized_linear_forward_global_scale_matches_manual_sidecar_path at ffi_tests.rs:2392, the only other ffi_tests case that runs native NVFP4; PR #1856's gate run recorded two nvfp4 group-16 aborts, consistent with both, and the mxfp4 cases compiled_qgelu_mlp_global_scale_non_nvfp4_multi_token_fallback_matches_reference and ..._no_sidecar_matches_unscaled) and every module sorting after ffi_tests, such as fused_moe_parity_tests, grouped_gemm_numeric_tests, hardware, layers, mla, paged_v2, rocm_arch, sampling* and speculative. Until this is fixed, anything that needs evidence from a later-sorting test must run that test directly (cargo test -p mlxcel-core --lib --features rocm -- <name>), as PRs #1883 and #2029 did.

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).

Update 2026-09-29: mxfp4 is no longer broken as far as the record shows. PR #1818 reports mxfp4 and mxfp8 quantized_matmul and gather_qmm matching a dequantized f32 reference (about 3e-3 relative for bf16), mxfp4 GPU quantize bit-identical to the CPU at 4096x4096 (the launch failure was arg_reduce grid geometry, patches-rocm/LOCAL_FIXES.md item 11), and gpt-oss-20b-MXFP4-Q4 generating coherent text. docs/installation.md:476 now tells users mxfp8 and mxfp4 run natively. #1808 still owns confirming that the original hang does not reproduce. NVFP4 is unchanged: no native kernel (docs/installation.md:477).

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. (Verified 2026-09-29: nvfp4_repack_strategy is still at :533-547; the affine branch is :851-869; the skip messages span :662-718, plus a separate affine-group skip at :864. The strategy's backend input is cfg!(feature = "cuda") in current_nvfp4_repack_strategy() (:551-557), a build flag, so a ROCm build takes DirectTranscode by default and produces the group-16 layout the ROCm dispatch above does not implement. There is also a third strategy, DenseNative, chosen when the dense override is set on CUDA or together with MLXCEL_NVFP4_NATIVE_REPACK. The load log line naming the chosen route already exists at :597-600, without a reason.)
• 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.

Where the table and the policy should live (added 2026-09-29)

Two anchors landed after this issue was written and are better than the "backend kind from #1803 or the build features" choice in step 1:

• Key the table on GpuBackendKind, read through mlxcel_core::hardware::gpu_backend_kind() (src/lib/mlxcel-core/src/hardware.rs:154-160 and :206, from PRs refactor(core): route custom kernels by GPU backend kind #1877 and feat(rocm): add a GPU vendor concept and AMD device reporting #1883), not on build features. The WebUI catalog already does exactly this in backend_status_for (src/server/webui/catalog_metadata.rs:337), whose comment records why: a cfg! chain sent every ROCm build to Unsupported (fix(rocm): the WebUI catalog reports every architecture unsupported on ROCm #1886) and cannot express a multi-backend build. GpuBackendKind is #[non_exhaustive], so a downstream match needs a wildcard arm; walk GpuBackendKind::ALL in a test so a new variant cannot silently fall into it. GpuBackendKind::None (a CPU-only build, or ROCm with no visible device) needs its own column.
• Put the load-time convert-or-reject hook in src/models/sanitize.rs, beside pre_ampere_cuda() (:1879) and bf16_to_f16_at_load() (:1803), the existing precedent for a backend-capability load policy, and route current_nvfp4_repack_strategy() through it instead of cfg!(feature = "cuda"). The table itself can sit in hardware.rs next to GpuBackendKind, since it describes the backend rather than a model.

Implementation plan

1. Add a backend quantization capability query (Rust side, fed by the backend kind from refactor(core): route custom kernels by GPU backend kind instead of treating every non-Metal GPU as CUDA #1803 or the build features): per mode, Native, ConvertTo(mode) or Unsupported. Initial ROCm table: affine native; mxfp8 native (after feat(rocm/quant): mxfp8 end-to-end on ROCm, including FP8 block checkpoints and the MoE gather path #1807 validates it); mxfp4 ConvertTo(affine 4-bit) until fix(rocm/quant): mxfp4 qmm hang and GPU quantize launch failure, with an affine 4-bit fallback #1808 fixes the kernels; NVFP4 ConvertTo(affine 4-bit). Metal and CUDA report exactly today's behavior. (2026-09-29: decided: use the backend kind, see the section above; the build-features option is rejected for the reason recorded there. The mxfp4 entry is stale: PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818's results support Native for mxfp4, with fix(rocm/quant): mxfp4 qmm hang and GPU quantize launch failure, with an affine 4-bit fallback #1808 confirming the hang is gone.)
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.
6. (Added 2026-09-29) Clear the gate: the two native-NVFP4 tests in ffi_tests.rs (:2392, :2980) must not abort on a backend without a native NVFP4 kernel. Gate them on the capability query from step 1 and skip with a printed reason, the pattern PR refactor(core): route custom kernels by GPU backend kind #1877 used for bitlinear_matmul's test, or make the path return Result. Then rerun the full -p mlxcel-core --lib target, since every test after the abort has been unobserved on ROCm and may surface new failures.

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.
• (Added 2026-09-29) make verify-rocm on gfx1151 finishes -p mlxcel-core --lib with no terminate called, and the run's full pass/fail list is recorded in the PR, because it covers tests no ROCm gate run has reached since the abort.
• (Added 2026-09-29) docs/installation.md:477 states the NVFP4 route on ROCm.

References

• src/models/sanitize.rs (nvfp4_repack_strategy 533-547, skip messages 662-709, bf16_to_f16_at_load 1803, pre_ampere_cuda 1879) (2026-09-29: skip messages now 662-718 and 864; the others are unchanged)
• src/models/fp8_block.rs (MXFP8_MODE line 68)
• Logit-trace method: docs/benchmarks.md, examples/logit_trace, scripts/compare_logit_traces.py
• (Added 2026-09-29) src/lib/mlxcel-core/src/hardware.rs (GpuBackendKind, gpu_backend_kind()); src/server/webui/catalog_metadata.rs:337 (backend_status_for, runtime-kind precedent); src/lib/mlxcel-core/src/ffi_tests.rs:2392, :2980 (native-NVFP4 tests); patches-rocm/mlx/backend/rocm/quantized/qmm.hip:4277-4293 (group-size dispatch)

Refresh log

2026-09-29

• Recorded that this is the only red item in the ROCm gate at d8d34e2b and therefore the effective blocker for feat(rocm/quant): mxfp8 end-to-end on ROCm, including FP8 block checkpoints and the MoE gather path #1807, fix(rocm/quant): mxfp4 qmm hang and GPU quantize launch failure, with an affine 4-bit fallback #1808 and test(rocm): correctness matrix against a Metal baseline and a verify-test-rocm gate #1809; named the aborting test and explained that the abort hides every later-sorting test, so later evidence must come from direct runs.
• Decided the anchor question left open in step 1: key the table on GpuBackendKind via gpu_backend_kind() (PRs refactor(core): route custom kernels by GPU backend kind #1877, feat(rocm): add a GPU vendor concept and AMD device reporting #1883), following the WebUI catalog's backend_status_for, and host the load policy beside pre_ampere_cuda() in sanitize.rs; noted that current_nvfp4_repack_strategy() still keys on cfg!(feature = "cuda").
• Marked the mxfp4 description and the initial mxfp4 table entry as stale per PR feat(rocm): vendor the ROCm backend into mlxcelverse, add rocm feature #1818's op-level and gpt-oss results.
• Corrected the sanitize.rs skip-message range and the affine-branch lines; added the DenseNative strategy and the existing route log line.
• Added a gate-clearing step and two acceptance criteria (no abort in -p mlxcel-core --lib, docs row).