[CUDA] gather_qmm ~24x slower than equal-FLOPs quantized_matmul on GB10 (sm_121) — makes quantized MoE prefill unusable

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)

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.

Start in backend/cuda/quantized/ and reproduce the reported gather_qmm and dense quantized_matmul timings with the benchmark in the issue on an sm_121 device. Compare the selected gather variants for the reported shapes and modes; done means a substantially faster gather path or documentation identifying which architectures use the fast path.