perf(rocm): add an mxfp4 arm to the warp-shared gather_qmv kernel for gpt-oss decode #2178

Part of #1814 (epic #1801)

Problem / Background

PR #2164 (#2106) sent sorted bf16 mxfp4 gather_qmm prefill to the expert-batched kernel, and gpt-oss-20b 512-token prefill on gfx1151 went from 7.53 to 526.3 tok/s. Decode did not move (median 8.45 tok/s off, 8.49 on), because a decode step never meets that kernel's gate (B = top_k = 4, below B >= 64). The guarded rocprofv3 --kernel-trace --stats run in docs/benchmark_results/rocm-moe-prefill-mxfp4-gfx1151-2026-10-06.md ("After the change") shows the per-row gather_qmv_kernel<hip_bfloat16, unsigned char, 4, 32, false> at 1.39 ms per call, 504 calls over seven decode steps (72 per step: 24 layers times gate, up and down), so about 100 ms of a roughly 115 ms step. TECHNICAL_REPORTS/2164-mxfp4-expert-batched-gather-qmm-20261007.en.md (section 7) names it the next target for gpt-oss decode. The trace itself is not committed.

Current Behavior

• GptOssSwitchGLU (src/models/gpt_oss.rs:529-560) sorts only when tokens * top_k >= 64, so decode calls gather_qmm unsorted with B = 4, M = 1, K = N = 2880, mxfp4 at group size 32, bf16 activations, no quantization biases.
• GatherQMM::eval_gpu (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip:6028) computes use_fast_gather_qmv (:6094-6100, true for transposed 4-bit), but every fast arm after it is affine-only: the wide gather (:6393), the idot gather and the warp-shared arm (:6461-6551, mode_ == Affine, group size 64). mxfp4 falls to the non-affine block at :6555, which launches gather_qmv_kernel<T, uint8_t, BITS, 32, false> (:5703, LOCAL_FIXES item 10) with 256-thread blocks and one thread per output column walking all of K serially.
• gather_qmv_warp_shared_kernel (:5058-5490) already has an AFFINE=false body (:5326-5480), but nothing instantiates it for gather. Its 4-bit path loops per group with k_local = lane * 8 and step THREADS_PER_COL * 8, so at group size 32 only lanes 0 to 3 of 16 do work in each group, and it decodes through the branchy fp4_e2m1_to_float switch (:1832).
• The same switch to the warp-shared kernel fixed f16 MoE decode before (LOCAL_FIXES item 12: Mixtral f16 went from about 205 ms per call in the per-row kernel to 1.44 ms). Items 8 and 10 are the E8M0 uint8_t scale dispatch this arm must keep.

Proposed Solution

1. In gather_qmv_warp_shared_kernel, add an if constexpr (!AFFINE && BITS == 4 && GROUP_SIZE == 32) path that replaces the per-group loop for that case. Each lane walks packed words across the whole shared_x chunk: for (int k = chunk_start + lane * 8; k < chunk_end; k += THREADS_PER_COL * 8), loads uint32_t packed = *reinterpret_cast<const uint32_t*>(&w_row[k / 2]), takes the group g = k / 32 (all eight nibbles of a word sit in one group, as in item 31), loads scale = load_scale_value<ScaleT, 32, false>(scales_row[g]), decodes the eight nibbles with fp4_e2m1_to_float_scaled (:1878), accumulates qx over the eight shared_x values with fmaf, then acc = fmaf(scale, qx * kFp4HalfBitsScale, acc). The end-of-kernel subgroup_reduce_sum_qmm<THREADS_PER_COL>(acc) and store stay as they are. Leave the existing AFFINE=false body in place for the other non-affine cases; nothing instantiates it after this change either.
2. In GatherQMM::eval_gpu, add an arm directly after the affine warp-shared arm (before #define has_bias at :6553): use_fast_gather_qmv && mode_ == QuantizationMode::Mxfp4 && x.dtype() == bfloat16 && group_size_ == 32 && bits_ == 4 && K % 32 == 0. It launches gather_qmv_warp_shared_kernel<hip_bfloat16, uint8_t, 4, 32, false, 16> when fast_threads_per_col == 16, else the WARP_SIZE instantiation, with fast_grid, fast_block, (const uint8_t*)scales_ptr, a null bias pointer, has_bias = false, use_sorted_rhs_schedule and implicit_x_batch_stride, then returns. That is two new instantiations, bf16 only, as item 31 limited itself to bf16 to hold qmm.hip compile time.
3. Update the comment at :6555 and LOCAL_FIXES item 10 to say bf16 mxfp4 at group size 32 now takes the warp-shared kernel, and add item 32 under "Fixes to the fork's kernels" naming qmm.hip, with the numbers below and the closing "Applies to the fork; kept in mlxcelverse under the 2026-10-06 fork policy, not proposed there." Fork policy: the fix stays in the mlxcelverse overlay as a LOCAL_FIXES entry; no docs/mlxcelverse/upstream/ package.

Rejected: instantiating the existing per-group AFFINE=false body as is (12 of 16 lanes idle per group at group size 32, plus the branchy decode); extending the expert-batched kernel to B = 4 (its block-per-expert grid launches E = 32 row blocks for four rows, and its gate is tuned for prefill).

Scope

In scope: qmm.hip (kernel path and dispatch arm), patches-rocm/LOCAL_FIXES.md (items 10 and 32), the new test below, and a results page under docs/benchmark_results/.

Out of scope: f16 mxfp4 and mxfp8 gather (keep the per-row kernel), the dense mxfp4 qmv_warp_shared_kernel (same per-group lane pattern; separate issue if profiling shows it), the expert-batched gate, Wave64 tuning.

Implementation Notes

• Switch for A/B: no new variable. MLX_ROCM_GATHER_QMV_USE_WARP=0 (read on every call, :6099) sets use_fast_gather_qmv = false, which sends mxfp4 back to the per-row kernel. Use it for the test and the measurement, and document it in the LOCAL_FIXES entry.
• Sorted calls: a sorted mxfp4 call that misses the expert-batched gate (prefill under 32 tokens, B / E < 4) now also takes the new arm, with use_sorted_rhs_schedule passed as the kernel's implicit_lhs argument exactly as the affine arm passes it. The test must cover it.
• Edge cases: rhs_idx >= E writes 0 (existing expert_valid handling); a partial last column block (col >= N) must not write; K not a multiple of 32 never reaches the arm (falls through to the per-row kernel). K = 2880 spans two shared_x chunks (2048 + 832), both multiples of 32.
• Error handling: no new failure paths. A launch error surfaces through the existing try_eval reporting (fix(rocm): surface ROCm GPU failures through Event::error instead of NaN or hangs #1804).

Acceptance Criteria

• Unsorted bf16 mxfp4 gather_qmm with B = 4 launches gather_qmv_warp_shared_kernel<hip_bfloat16, unsigned char, 4, 32, false, 16> and no gather_qmv_kernel, shown by a guarded rocprofv3 --kernel-trace --stats run of gpt-oss-20b decode, with its per-call time recorded.
• New test mxfp4_warp_shared_gather_qmv_matches_per_row_and_reference in tests/rocm_mxfp4_quant.rs (reusing its quantize_on, dequantize_cpu and gather_reference): 32 experts, top 4, K = 2880, N = 2880 and N = 516, bf16, unsorted with 1 and 8 tokens and sorted with 16 tokens (B = 64, misses the expert-batched gate). The default path's relative error against the dequantized f32 reference stays within the per-row path's (MLX_ROCM_GATHER_QMV_USE_WARP=0) error times 1.05 and under 2e-2; two default runs are bit-identical; at N = K = 2880 the default and per-row outputs differ in at least one byte (proves dispatch). The test sets the variable under lock_default_device and restores it. Dropping the sign bit in the new decode must fail it (record the value in the PR).
• gather_qmm_matches_per_expert_reference and the expert-batched tests still pass.
• Six alternated rounds, MLX_ROCM_GATHER_QMV_USE_WARP=0 against default, every run under scripts/rocm_gpu_guard.sh: decode median improves by more than the about 1% host spread perf(rocm): add a batched gather_qmm path for mxfp4 MoE prefill #2106 measured, and 512-token prefill stays within that spread. If decode does not improve, the arm is not merged as default; record the numbers in the issue instead.
• Teacher-forced examples/logit_trace over tests/fixtures/wikitext2_excerpt.txt (w8 = 8 80 8 512), per-row against default, scripts/compare_logit_traces.py --decided 2.0: 0 decided mismatches.
• Results page docs/benchmark_results/rocm-moe-decode-mxfp4-gfx1151-<date>.md and LOCAL_FIXES item 32 recorded; the gpt-oss decode path in a real mlxcel run takes the new arm (integrated, no opt-in needed).

Verification

cargo build --release --features rocm
cargo test --release --features rocm --test rocm_mxfp4_quant -- --test-threads=1
cargo test --release --features rocm --test rocm_gather_qmm_expert_batched -- --test-threads=1
make verify-rocm-overlay
# per round, alternate the two arms
scripts/rocm_gpu_guard.sh -- env MLX_ROCM_GATHER_QMV_USE_WARP=0 ./target/release/mlxcel-bench-decode -m models/mlx/gpt-oss-20b-MXFP4-Q4 -p x --prompt-tokens 512 -n 128 --warmup-tokens 20 --ignore-eos
scripts/rocm_gpu_guard.sh -- ./target/release/mlxcel-bench-decode -m models/mlx/gpt-oss-20b-MXFP4-Q4 -p x --prompt-tokens 512 -n 128 --warmup-tokens 20 --ignore-eos
scripts/rocm_gpu_guard.sh -- rocprofv3 --kernel-trace --stats -f csv -d prof -o run -- ./target/release/mlxcel-bench-decode -m models/mlx/gpt-oss-20b-MXFP4-Q4 -p x --prompt-tokens 512 -n 8 --warmup-tokens 1 --ignore-eos

A pass: tests green, the trace shows the warp-shared instantiation instead of the per-row one, and the decode median beats the per-row arm outside the host spread.

Implemented in #2235 (LOCAL_FIXES item 41, since 32 and 40 were taken). Decode median 8.32 to 63.62 tok/s over six alternated guarded rounds, prefill unchanged; see docs/benchmark_results/rocm-moe-decode-mxfp4-gfx1151-2026-10-08.md.