Skip to content

perf(rocm): add a batched gather_qmm path for mxfp4 MoE prefill #2106

Closed

Description

@inureyes

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_qmm against 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

  1. added
    area:coremlxcel-core: MLX FFI, primitives, KV cache, layers
    platform:linuxLinux (CUDA / packaging) specific
    on Oct 5, 2026
  2. added a commit that references this issue on Oct 6, 2026
    08e2269
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:coremlxcel-core: MLX FFI, primitives, KV cache, layersplatform:linuxLinux (CUDA / packaging) specificpriority:mediumMedium prioritystatus:doneCompletedtype:performancePerformance improvements

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions