Repository navigation
[CUDA] gather_qmm ~24x slower than equal-FLOPs quantized_matmul on GB10 (sm_121) — makes quantized MoE prefill unusable #4339
Description
Describe the bug
On an NVIDIA GB10 (DGX Spark class, sm_121), mx.gather_qmm runs at ~0.85
TFLOPs for MoE-shaped workloads while a dense mx.quantized_matmul doing the
same total FLOPs reaches ~20.3 TFLOPs on the same device — a ~24x gap.
sorted_indices=True makes no measurable difference.
The practical effect: prompt processing for every quantized MoE model is an
order of magnitude slower than the same model dequantized to bf16, and slower
than much larger dense quantized models.
Model-level numbers (mlx_lm generate, ~1.9k-token prompt, single stream)
| Model | Format | Prefill tok/s |
|---|---|---|
| Llama-3.1-8B (dense) | bf16 | 1,631 |
| Llama-3.1-8B (dense) | affine 4-bit | 1,378 |
| gpt-oss-20b (MoE) | bf16 (dequantized) | 860 |
| gpt-oss-20b (MoE) | MXFP4-Q8 | 132 |
| Qwen3-Next-80B-A3B (MoE) | affine 4-bit | 105 |
| Qwen3.5-122B-A10B (MoE) | affine 4-bit | 44 |
Dense quantized is ~85% of bf16 — the plain qmm path is healthy on this
device. The same MoE model is 6.5x faster with bf16 weights than with
quantized weights, isolating the regression to the quantized gather path.
Kernel microbenchmark (shapes copied from mlx_lm QuantizedSwitchLinear)
import time
import mlx.core as mx
T, A, H, F, E = 2048, 8, 2048, 768, 64 # tokens, active experts, hidden, ffn, experts
x = mx.random.normal((T, A, 1, H)).astype(mx.bfloat16)
w = mx.random.normal((E, F, H)).astype(mx.bfloat16)
wq, sc, bs = mx.quantize(w, group_size=64, bits=4)
idx = mx.random.randint(0, E, (T, A)).astype(mx.uint32)
flops = 2 * T * A * H * F
def bench(fn, label, reps=10):
for _ in range(3):
mx.eval(fn())
mx.synchronize()
t0 = time.perf_counter()
for _ in range(reps):
mx.eval(fn())
mx.synchronize()
dt = (time.perf_counter() - t0) / reps
print(f"{label}: {dt*1000:8.2f} ms {flops/dt/1e12:6.2f} TFLOPs")
bench(lambda: mx.gather_qmm(x, wq, sc, bs, rhs_indices=idx, transpose=True,
group_size=64, bits=4, mode="affine",
sorted_indices=False), "gather_qmm unsorted")
bench(lambda: mx.gather_qmm(x, wq, sc, bs, rhs_indices=idx, transpose=True,
group_size=64, bits=4, mode="affine",
sorted_indices=True), "gather_qmm sorted ")
xd = mx.random.normal((T * A, H)).astype(mx.bfloat16)
wd, sd, bd = mx.quantize(mx.random.normal((F, H)).astype(mx.bfloat16),
group_size=64, bits=4)
bench(lambda: mx.quantized_matmul(xd, wd, sd, bd, transpose=True,
group_size=64, bits=4), "dense qmm (equal FLOPs)")Output on GB10:
gather_qmm unsorted: 60.33 ms 0.85 TFLOPs
gather_qmm sorted : 60.16 ms 0.86 TFLOPs
dense qmm (equal FLOPs): 2.54 ms 20.33 TFLOPs
For reference, plain GEMM on the same device: bf16 mx.matmul 64 TFLOPs,
affine 4-bit mx.quantized_matmul 48 TFLOPs (2048x4096x14336).
Expected behavior
gather_qmm within striking distance of dense qmm for large-token MoE prompt
processing (as it is on Metal, where quantized MoE prefill is fast), or at
least documentation of which architectures take the fast path. Reading
backend/cuda/quantized/, GatherQMM appears to have only sm80/naive variants
(no sm90 tier, nothing Blackwell-aware); it looks like sm_121 with these
shapes/modes ends up on the naive path.
Desktop
- Device: NVIDIA GB10 (
sm_121, compute capability 12.1), 128 GB unified - OS: Linux 6.17.0-1029-nvidia aarch64 (DGX OS)
- Driver 580.173.02, CUDA 13.0
- mlx 0.32.0 / mlx-cuda-13 0.32.0 / mlx-lm 0.31.3 (uv venv, Python 3.13)
Happy to run patches or nightly builds on this hardware — the box is available
for benchmarking.
Activity
NORDIT commentedon Aug 18, 2026
Follow-up with bf16 MoE data on the same GB10 — it shows the gather bottleneck is not specific to quantized weights, and that it scales with expert granularity.
Same-device prefill, mlx_lm generate, ~1.9k-token prompt:
Dense models on the same GB10 for contrast: Llama-3.1-8B bf16 1,631 tok/s, affine 4-bit 1,378 tok/s.
So on CUDA both
gather_qmmand the bf16 MoE gather path fall far below the dense kernels, and the fine-grained-expert models (Qwen A3B family) are hit hardest — 6.5x worse than a coarse-expert MoE of similar active size in bf16, and ~7-15x worse than the same file on Metal. In a distributed prefill/decode setup (exo) this currently inverts the expected prefill advantage of the GB10 over an M3 Ultra: end-to-end, a ~7k-token request is 5.4 s Mac-only vs 30.7 s with GB10 prefill (bf16 35B, identical safetensors both sides, verified by checksum).Happy to profile or test candidate kernels on this box.
manunicholasjacob commentedon Sep 2, 2026
Same microbenchmark on sm_86 (RTX 3050 Laptop, 4 GB, driver 596.08 under WSL2, mlx 0.32.2 + mlx-cuda-13 0.32.2), your shapes (T=2048, A=8, H=2048, F=768, E=64, affine 4-bit g64). Interleaved, six rounds in rotated order, 3 warmups and 10 timed reps with
mx.synchronize()per rep; medians across rounds, spread in parentheses. A 4 GB laptop part is not comparable to a GB10 in absolute terms, only the ratios are.So 28x here against your 24x, and it is structural rather than Blackwell-specific:
supports_qmm_sm80only checks compute capability major >= 8, so sm_86 and sm_121 both runqmm_sm80_kernel, which for the (T, A, 1, H) input has m = 1, takes a 16-row CTA tile and launches a (1, 6, 16384) grid, one CTA per token-expert pair with 15 of 16 rows empty and the expert's 4-bit slab dequantised once per row. The bf16 path already has the tier that is missing: with tokens sorted by expert andsorted_indices=True,GatherMMdispatchescutlass_grouped_gemm_unalignedand lands within about 2x of dense, whilegather_qmmnever consults the flag on this path and only picks up L2 locality from sorted indices. #4390 is building a proper grouped GEMM for the bf16 case on top of that; the quantized path would need the same tier, since the 4-bit slab is what makes the per-row dequantise expensive. Scripts and raw JSONL available if useful.arcusbuilds commentedon Oct 10, 2026
I'd like to take the affine case. Since #4554, sorted
gather_qmmon CUDA usesgather_qmm_rhs_sm80, butquantized.cppskips it for affine (mode_ != QuantizationMode::Affine), so affine 4-bit still runs the slowqmm_sm80path.Plan: add bias support to the rhs sm80 dequant, add affine sorted tests to
test_quantized.py, and post before/after numbers from the script above on sm_89.@nastya236 does #4390 cover any of this, or is it bf16 only?
nastya236 commentedon Oct 11, 2026
Hey! @mattt @arcusbuilds @manunicholasjacob @NORDIT thanks for looking into this. Current gather_qmm is written for sm80 so it is expected that it is slow on sm120. #4390 covers only bf16, but once it is merge, I will do a pr for
grouped_qmm. I will come back to grouped_mm pull request as soon as possible and we will close this issue. For cuda we won't support affine quantization, onlymxfp8andnvfp4.