[inureyes]

Summary

ROCm had no NVFP4 kernel, so ModelOpt NVFP4 checkpoints were converted to affine 4-bit at load and MLX-native NVFP4 planes were refused (#1806). This adds native NVFP4 (E2M1 values, E4M3 scales per 16, optional float32 global scale per expert) to QuantizedMatmul and GatherQMM and marks (Rocm, Nvfp4) Native.

• New overlay file quantized/nvfp4_qmv.hip (included by qmm.hip): nvfp4_qmv_kernel computes eight rows per block (one for M = 1 and unsorted gather decode) and reloads a weight group only when the expert changes, so dense multi-row calls (speculative-decode verify widths) and sorted gather prefill read each group once per eight rows. It serves QuantizedMatmul below the dequantize + GEMM crossover and every NVFP4 GatherQMM call, applying global_scale[expert] / (448 * 6) once per output. nvfp4_dequantize_kernel expands a plane exactly; select_qmm_route admits NVFP4 to the dequantize + hipBLASLt/rocBLAS route, which covers prefill-size and non-transposed GEMMs.
• qdequant.hpp: branch-free E4M3 decode beside the moved fp4 decode (both exact, power-of-two offsets folded into the final multiply).
• Wave64: the kernels use 16-lane column groups with width-16 shuffles and never read WARP_SIZE, so they are lane-free and need no wavefront gate (not testable on this host).
• The NVFP4 gather path returns before the mxfp4/mxfp8 group-32 check. QuantizedMatmul now refuses an unexpected fourth input in a non-affine mode instead of ignoring it.
• validate_native_global_scale (switch_layers) admitted only Metal and CUDA; it now admits ROCm. Sanitize route tests, layers route-line tests and docs (installation ROCm table, env vars, supported models) follow the table change. LOCAL_FIXES.md item 63 (rebased onto fix(rocm): split per-row softmax and norm grids past the HIP limit #2627, which took 62).

Numerical results (gfx1151, vs f32 x @ dequantize(w).T on CPU)

Max abs error / error relative to max |reference|, worst case per path:

Path	f32	f16	bf16
qmm qmv M=1	7.6e-5 / 5.2e-7	5.9e-2 / 4.0e-4	4.7e-1 / 3.1e-3
qmm multi-row qmv M=6,16,33	9.5e-5 / 6.1e-7	6.2e-2 / 3.6e-4	5.0e-1 / 2.9e-3
qmm dequant+GEMM M=64,300	4.7e-4 / 2.3e-6	6.3e-2 / 4.1e-4	5.0e-1 / 3.3e-3
qmm batched w (qmv M=3 / GEMM M=64)	1.1e-5 / 1.4e-7, 9.9e-5 / 1.0e-6	n/a	2.5e-1 / 2.9e-3, 2.5e-1 / 2.4e-3
qmm non-transposed (GEMM)	5.7e-5 / 6.1e-7	n/a	2.4e-1 / 2.7e-3
qmm + global scale after matmul	1.5e-4 / 1.2e-6	6.8e-2 / 4.9e-4	5.3e-1 / 5.4e-3
gather unsorted decode, no gs	2.5e-5 / 2.2e-7	3.1e-2 / 3.0e-4	4.0e-1 / 2.7e-3
gather unsorted decode, per-expert gs	2.3e-5 / 2.2e-7	3.1e-2 / 2.5e-4	2.5e-1 / 2.2e-3
gather sorted prefill (B=128), no gs	3.1e-5 / 2.3e-7	5.9e-2 / 4.4e-4	3.0e-1 / 2.2e-3
gather sorted prefill, per-expert gs	3.1e-5 / 2.3e-7	3.1e-2 / 2.0e-4	3.0e-1 / 2.0e-3
gather M=3 rows per expert, no gs / gs	3.1e-5 / 2.2e-7, 2.3e-5 / 1.8e-7	3.1e-2 / 2.4e-4, 3.1e-2 / 2.5e-4	2.5e-1 / 2.1e-3, 3.9e-1 / 2.8e-3

f16/bf16 errors are output rounding; f32 matches to accumulation order.

Test plan

• tests/rocm_nvfp4_quant.rs and tests/rocm_nvfp4_quant/gather.rs (new): 7 tests pass on gfx1151 under the GPU guard after the post-rebase overlay rebuild (worst relative error f32 2.1e-6, f16 4.9e-4, bf16 5.4e-3)
• The six formerly gated tests take the native branch and pass: the three switch_layers::global_scale_tests cases, gemma_unified::...unified_nvfp4_fused_experts_keep_their_native_global_scale, mistral4::modelopt_stack_tests::..., phimoe::sanitize_tests::...; also the ROCm-skipped sanitize_tests NVFP4 loads and ffi_tests native-NVFP4 global-scale tests (post-rebase rerun: mlxcel lib filters 56 passed, mlxcel-core 11 passed)
• Regression: rocm_mxfp4_quant, rocm_hip_grid_clamp gather_qmv, rocm_gather_qmm_expert_batched
• make verify-llama-compat verify-env-var-docs verify-kernel-port-dispatch verify-kernel-dtype-keys verify-rocm-overlay
• cargo clippy --profile test-fast --features rocm -p mlxcel -p mlxcel-core --lib --tests -- -D warnings
• Real-checkpoint smoke: not run. No NVFP4 checkpoint is in the local HF cache and large downloads were out of scope on this host.

Changes during review

• The K-multiple-of-16 check now applies to the quantized axis (transpose_ ? K : N), so a non-transposed weight with K = 17 reaches the dequantize + GEMM route instead of being refused.
• nvfp4_qmv caps grid x at UINT32_MAX / 16 (blockDim.x is 16) with a clear error.
• check_nvfp4_planes requires scales.shape == w.shape[:-1] + [w.shape[-1] / 2] with one-byte scales in both the QuantizedMatmul and GatherQMM arms, and the dequantize sizes its grid from the weight. Before this, a malformed scale plane (which MLX's validate_quantized_input lets through) could make the dequantize kernel write past its output and the kernels read past the scales.
• New tests: a partial 16-column block (N = 272), out-of-range expert indices (E, E+1, i32::MAX) that must produce exact zeros, an untransposed K = 17 case, a mismatched scale plane that must be refused, and assert_reconstructs, which ties the global-scale reference to the source weights rather than to the kernel's constant.
• The global-scale factor global_scale / (448 * 6) was checked against MLX's fp_quantize/fp_dequantize in mlx/ops.cpp and the CUDA fp_gather_qmv.

Follow-ups: an expert-batched or grouped dequantize + GEMM path for NVFP4 gather prefill (it currently reads each expert's weights once per eight sorted rows); wave64 not measured; nvfp4_qmv_kernel does not bound lhs_indices against the x batch (same as most existing gather kernels; mlxcel generates them); sorted NVFP4 gather prefill may be slower than the former affine conversion with its expert-batched routes (not measured, no checkpoint); GatherQMM::eval_gpu divides by M*N before an empty-output check (pre-existing).

Closes #2617

[inureyes]

Implementation Review Summary

Intent

Native NVFP4 (E2M1, E4M3 scales per 16, optional per-expert global scale) in ROCm QuantizedMatmul and GatherQMM, with NVFP4 marked Native on ROCm.

Findings Addressed

• Non-transposed NVFP4 checked K instead of the quantized axis N, refusing valid shapes such as x [1, 17] against a [17, 32] weight (MEDIUM)
• nvfp4_qmv grid x cap was INT32_MAX; HIP's limit is UINT32_MAX / 16 blocks at 16 threads in x (LOW)
• Test gaps: N not a multiple of 16, out-of-range expert indices, and a reference check independent of the kernel's global-scale formula (MEDIUM)

Remaining Items

• Real NVFP4 checkpoint smoke (none cached; out of scope by disk policy)
• Wave64 not run on hardware; the kernels are lane-free by construction (16-lane groups, width-16 shuffles, no barriers, no WARP_SIZE)

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into project code flow
• Project conventions followed
• Existing modules reused where applicable
• No unintended structural changes
• Tests pass

The global-scale convention matches MLX: fp_quantize encodes block scales times 448 * 6 / global_scale and fp_dequantize and the CUDA fp_gather_qmv multiply by global_scale / (448 * 6). Post-rebase run on gfx1151 (onto 232499d): rocm_nvfp4_quant 6/6 (max relative error f32 2.1e-6, f16 4.9e-4, bf16 5.4e-3), rocm_mxfp4_quant 5/5, rocm_hip_grid_clamp gather_qmv 4/4, and 56 mlxcel plus 11 mlxcel-core NVFP4 and global-scale lib tests, including the six formerly gated tests on the native branch.

[inureyes]

Security and performance review

Fixed in f2acc66:

• HIGH: an oversized NVFP4 scale plane made nvfp4_dequantize_kernel write past the dequantized weight (grid sized from scales.size()), and a short or extra-axis plane made every NVFP4 kernel read past the scales. MLX's op check compares only the batch prefix and last axis, so a malformed checkpoint reaches the backend. check_nvfp4_planes now requires scales.shape == w.shape[:-1] + [w.shape[-1] / 2] in both arms, and the dequantize sizes from the weight. New refusal test covers both routes.

Not fixed (below the auto-fix threshold):

• MEDIUM: nvfp4_qmv_kernel does not bound lhs_indices against the x batch count, so an out-of-range lhs index reads past x. Matches most existing gather kernels; the tiled ones clamp it. mlxcel builds lhs indices internally.
• MEDIUM (perf): sorted NVFP4 gather prefill reads each expert's weights once per eight rows, while the affine conversion it replaces had expert-batched GEMM routes. Unmeasured without a checkpoint; already listed as a follow-up.

Checked without findings: global-scale length vs expert count, rhs bounds (zeroed rows), partial row and column tiles, 64-bit index math, grid limits, and the dequantize + GEMM plane (one weight matrix, same bound and LRU cap as affine). Affine, mxfp4 and fp8 routing in select_qmm_route is unchanged.

Verification on gfx1151: rocm_nvfp4_quant 7/7, rocm_mxfp4_quant 5/5, the five contract verify targets, clippy on the test target.