perf(rocm): give the dense mxfp4 qmv kernels the warp-shared word path #2246

Part of #1801

Problem / Background

PR #2235 (#2178) gave gather_qmv_warp_shared_kernel an mxfp4 word path: each lane takes one packed 32-bit word (eight e2m1 nibbles, all inside one 32-value group) per step across the whole shared_x chunk, decodes without a branch through fp4_e2m1_to_float_scaled, and scales the word's dot product once. gpt-oss-20b decode went from 8.32 to 63.62 tok/s. Its results page (docs/benchmark_results/rocm-moe-decode-mxfp4-gfx1151-2026-10-08.md:63) notes that the dense mxfp4 kernels were not changed. Whether dense mxfp4 GEMVs are a meaningful share of any decode is not measured.

Current Behavior

In src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip (origin/main 1b4e3657), the PR #2235 page says the dense mxfp4 kernel has the same per-group lane pattern. That holds for the batched kernel only (AFFINE == false, BITS == 4 code below):

• qmv_warp_shared_batched_kernel (:2361, used when can_use_batched_qmv, mxfp4 loop at :2627-2660): same pattern as the old gather kernel. Lanes start at lane * 8 inside each group, so at group size 32 only lanes 0 to 3 of 16 do work, and every nibble goes through dequantize_value<4, false>, which reaches the branchy fp4_e2m1_to_float switch (:1892, :2003).
• qmv_warp_shared_kernel (:2015, non-batched transposed qmv, mxfp4 loop at :2259-2270): all 16 lanes work (stride THREADS_PER_COL, two values per lane per group), but each value is a separate byte load and shift (unpack_packed_value_fast, :1858) followed by the same branchy decode. So the problem here is the per-nibble decode, not idle lanes. qmv_warp_noshared_kernel (:2766, mxfp4 loop at :2904) does the same per-element loop.

The dispatch for these is the non-affine branch of QuantizedMatmul::eval_gpu (:4313 onwards, fast_threads_per_col 16 or WARP_SIZE).

Proposed Solution

1. Profile first. Take a guarded rocprofv3 --kernel-trace --stats decode trace (scripts/rocm_decode_profile.sh or the perf(rocm): add an mxfp4 arm to the warp-shared gather_qmv kernel #2235 recipe) of gpt-oss-20b and of any other mxfp4 checkpoint on the host, and report the share of decode GPU time in qmv_warp_shared{,_batched}_kernel<*, unsigned char, 4, 32, false, *> and qmv_warp_noshared_kernel<..., 4, 32, false, ...>. If the combined share is under 5%, record the numbers in this issue and close it without code changes.
2. If it is 5% or more, add the kMxfp4WordPath (!AFFINE && BITS == 4 && GROUP_SIZE == 32) block from gather_qmv_warp_shared_kernel (:5217-5247) to each kernel the trace shows, skipping its per-group loop for that case. K % 32 == 0 holds for every mxfp4 weight (group size 32), and CHUNK_SIZE (2048) is a multiple of 32, so no word straddles a chunk. Keep MLX_ROCM_QMV_USE_WARP=0 as the way back to the old kernel. Add a patches-rocm/LOCAL_FIXES.md entry (mlxcelverse only).

Acceptance Criteria

• The profile numbers are posted on this issue, whichever branch is taken.
• If implemented: a test in tests/rocm_mxfp4_quant.rs, modeled on mxfp4_warp_shared_gather_qmv_matches_per_row_and_reference, checks dense mxfp4 quantized_matmul (one row, and a batched case) against a CPU-dequantized f32 reference with relative L2 error no worse than the old kernel's, and runs the old kernel via MLX_ROCM_QMV_USE_WARP=0 for comparison; a sign-bit mutation must fail it.
• If implemented: guarded decode for the profiled checkpoint, at least three alternated rounds per arm, ranges reported; logit_trace default vs old kernel, 0 decided mismatches.

Verification

scripts/rocm_gpu_guard.sh -- scripts/rocm_decode_profile.sh   # see docs/benchmarks.md for arguments
cargo test --profile test-fast --features rocm --test rocm_mxfp4_quant -- --test-threads=1
make verify-rocm

Related: #2178, PR #2235, TECHNICAL_REPORTS/2235-rocm-mxfp4-warp-shared-gather-qmv-20261008.en.md, #2061, #1814.