Repository navigation
perf(rocm): add a batched gather_qmm path for mxfp4 MoE prefill #2106
Description
Part of #1814 (epic #1801). Split out of #2066, which fixed the affine expert-batched gather_qmm kernel and, per its step 5, files the mxfp4 case separately.
Problem
On gfx1151 (Radeon 8060S), measured 2026-10-05 at main c05d543 plus the #2066 branch (which does not touch this path), gpt-oss-20b-MXFP4-Q4 prefills 512 tokens in about 67.8 s (7.55 tok/s). A rocprofv3 --kernel-trace --stats profile of the prefill puts 99.78% of GPU time (67.7 s over 72 calls, about 940 ms per call) in gather_qmv_kernel<hip_bfloat16, unsigned char, 4, 32, false>, the generic per-row non-affine gather kernel that patches-rocm/LOCAL_FIXES.md item 10 routed mxfp4/mxfp8 to.
Cause
GatherQMM::eval_gpu in src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip gates every faster gather path (expert-batched, wide, tiled, warp-shared) on mode_ == QuantizationMode::Affine, so mxfp4 runs one thread block row per batch element and rereads each expert's weights for every routed row. gpt-oss has 32 experts and top-4 routing, so a 512-token prefill gives B = 2048 sorted rows, about 64 rows per expert.
Proposed Direction
Give mxfp4 (E8M0 uint8 scales, group size 32, no biases) a sorted expert-batched path like the affine gather_qmv_expert_batched_kernel, which reads each expert's weights once per 4 rows and reads the index arrays through the batch strides. An mxfp4 arm of the warp-shared kernel is an acceptable first step. Respect the ahead-of-time compile cost noted in LOCAL_FIXES item 17: add only the instantiations needed.
Acceptance Criteria
- A test compares sorted mxfp4
gather_qmmagainst the unsorted path and a dequantized f32 reference at the gpt-oss shape (K = 2880, E = 32, top 4). - 512-token prefill tok/s for gpt-oss-20b is reported before and after, measured under
scripts/rocm_gpu_guard.sh, medians of three. - Decode throughput is unchanged.
-
The LOCAL_FIXES entry is marked as an upstreaming candidate for chore(rocm): mlxcelverse ROCm fork sync script, MLX pin-bump procedure, and upstreaming local fixes #1813.Superseded by the maintainer's 2026-10-06 fork policy (ROCm fork fixes stay in mlxcelverse): the change is recorded as LOCAL_FIXES item 31 without the upstreaming mark.
Reference
Profile and numbers: docs/benchmark_results/rocm-moe-prefill-gfx1151-2026-10-05.md (added by the #2066 PR).
Implemented in #2164 (prefill 7.53 to 526.3 tok/s median on gfx1151, decode unchanged).
Activity
fix(rocm): fix and enable the expert-batched gather_qmm prefill kernel
fix(rocm): fix and enable the expert-batched gather_qmm prefill kernel (
test(rocm): check that sorted mxfp4 gather_qmm reaches the batched ke…
docs: add technical report for PR #2164
perf(rocm): batch sorted mxfp4 gather_qmm prefill by expert (#2164)