perf(rocm): Expert-batched NVFP4 sorted-gather prefill, validated on a real checkpoint #2665

Problem

On ROCm, GatherQMM::eval_gpu sends every NVFP4 gather to nvfp4_qmv (overlays/rocm/mlx/backend/rocm/quantized/qmm.hip:6237-6281, kernel in quantized/nvfp4_qmv.hip), sorted prefill included. With right_sorted_ && B >= 64 it shares one weight-group load across 8 rows (multi_row), so a 512-token prefill still reads each routed expert's weights once per 8 rows. Affine and mxfp4 sorted gathers get gather_qmv_expert_batched_kernel instead (qmm.hip:6462-6481, which finds each expert's run with find_sorted_expert_run at :4689), measured at 1.7x to 70x prefill on gfx1151 (LOCAL_FIXES.md items 9 and 31). Native NVFP4 (#2617, #2629) has also never run on a real checkpoint on ROCm: docs/installation.md:1118 says so, and only the synthetic tests in tests/rocm_nvfp4_quant/gather.rs cover it. None of the NVFP4 checkpoints are in ~/.cache/huggingface/hub.

Plan

Target checkpoint: mlx-community/Nemotron-Cascade-2-30B-A3B-nvfp4 (17.8 GB, nemotron_h, which mlxcel supports, 128 experts, top-6). Fallback: mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 (19.3 GB). Free disk is about 72 GB.

Add a sorted-prefill route for NVFP4 under the same gate shape as the affine one (transpose_ && right_sorted_ && M == 1 && B >= 64). Build it either as an NVFP4 instantiation of the expert-batched kernel or as a per-run nvfp4_dequantize_dense plus dequant_rocblas_gemm (both already in qmm.hip); pick the faster on the checkpoint. Do not copy the affine route's E <= 64 limit unless a measurement justifies it, since the target has 128 experts. Apply the per-expert global scale as nvfp4_qmv does, zero rows routed past the last expert, and honor MLX_ROCM_GATHER_QMV_EXPERT_BATCHED=0.

Implement everything in one PR with one verification pass at the end and as few pushes as possible.

Acceptance Criteria

• Baseline: the checkpoint loads and generates coherent greedy text on gfx1151, with prefill tok/s (512-token prompt) and decode tok/s from scripts/bench_decode.sh --prompt-tokens 512 recorded before any code change
• New tests in tests/rocm_nvfp4_quant/gather.rs with sorted B large enough to take the route and E = 128, with and without global scale, f32/f16/bf16, match the CPU reference at tolerance() (tests/rocm_nvfp4_quant.rs:142), and fail when the route's global scale or expert offset is broken on purpose
• Measured prefill speedup on the checkpoint with byte-identical greedy output, or measurements showing the existing path is already as fast, in which case record them and keep the code
• A new LOCAL_FIXES.md item (the next free number at merge time; 67 is the highest today and test(rocm): Cover donated-cotangent LayerNorm VJP, gemv fallback, empty norms; fix item 46 #2663 also adds one) records the route and the numbers; docs/installation.md:1118 names the validated checkpoint

Verification

scripts/rocm_gpu_guard.sh --hold -- cargo test --profile test-fast --features rocm --test rocm_nvfp4_quant -- --test-threads=1
scripts/rocm_gpu_guard.sh --hold -- ./scripts/bench_decode.sh <checkpoint dir> --prompt-tokens 512
make verify-rocm