Use matrix kernels for global-scale gather_qqmm

benchmarks/python/gather_qqmm_bench.py

@@ -0,0 +1,108 @@
+# Copyright © 2026 Apple Inc.
+
+import argparse
+
+import mlx.core as mx
+from time_utils import time_fn
+
+SHAPES = {
+    "qwen3.6-35b-a3b-up": {
+        "experts": 256,
+        "top_k": 8,
+        "input": 2048,
+        "output": 512,
+    },
+    "qwen3.6-35b-a3b-down": {
+        "experts": 256,
+        "top_k": 8,
+        "input": 512,
+        "output": 2048,
+    },
+}
+
+
+def make_weights(shape):
+    weight = mx.random.uniform(
+        low=-0.5,
+        high=0.5,
+        shape=(shape["experts"], shape["output"], shape["input"]),
+    ).astype(mx.bfloat16)
+    quantized, scales = mx.quantize(
+        weight,
+        mode="nvfp4",
+        global_scale=mx.array(1.0, dtype=mx.float32),
+    )
+    global_scale = mx.ones((shape["experts"],), dtype=mx.float32)
+    mx.eval(quantized, scales, global_scale)
+    return quantized, scales, global_scale
+
+
+def gather_qqmm(x, weight, scales, global_scale_x, global_scale_w, rhs, sorted_):
+    return mx.gather_qqmm(
+        x,
+        weight,
+        scales,
+        rhs_indices=rhs,
+        mode="nvfp4",
+        global_scale_x=global_scale_x,
+        global_scale_w=global_scale_w,
+        sorted_indices=sorted_,
+    )
+
+
+def benchmark(shape, workload, tokens):
+    weight, scales, global_scale_w = make_weights(shape)
+    global_scale_x = mx.array(1.0, dtype=mx.float32)
+
+    if workload == "prompt":
+        routes = tokens * shape["top_k"]
+        x = mx.random.uniform(
+            low=-0.5, high=0.5, shape=(routes, 1, shape["input"])
+        ).astype(mx.bfloat16)
+        rhs = mx.repeat(
+            mx.arange(shape["experts"], dtype=mx.uint32),
+            (routes + shape["experts"] - 1) // shape["experts"],
+        )[:routes]
+        sorted_ = True
+    elif workload == "matrix":
+        x = mx.random.uniform(
+            low=-0.5,
+            high=0.5,
+            shape=(shape["top_k"], tokens, shape["input"]),
+        ).astype(mx.bfloat16)
+        rhs = (mx.arange(shape["top_k"], dtype=mx.uint32) * 17) % shape["experts"]
+        sorted_ = False
+    else:
+        x = mx.random.uniform(
+            low=-0.5, high=0.5, shape=(1, 1, 1, shape["input"])
+        ).astype(mx.bfloat16)
+        rhs = (
+            (mx.arange(shape["top_k"], dtype=mx.uint32) * 17) % shape["experts"]
+        ).reshape(1, shape["top_k"])
+        sorted_ = False
+
+    mx.eval(x, rhs)
+    time_fn(
+        gather_qqmm,
+        x,
+        weight,
+        scales,
+        global_scale_x,
+        global_scale_w,
+        rhs,
+        sorted_,
+        msg=(
+            f"{workload} routes={rhs.size} " f"N={shape['output']} K={shape['input']}"
+        ),
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--shape", choices=SHAPES, default="qwen3.6-35b-a3b-up")
+    parser.add_argument(
+        "--workload", choices=("prompt", "matrix", "decode"), default="prompt"
+    )
+    parser.add_argument("--tokens", type=int, default=2048)
+    args = parser.parse_args()
+    benchmark(SHAPES[args.shape], args.workload, args.tokens)

mlx/backend/metal/quantized.cpp

@@ -2148,6 +2148,57 @@ void GatherQQMM::eval_gpu(const std::vector<array>& inputs, array& out) {
     global_scale_w = ensure_row_contiguous(gs_e, d, s);
   }
 
+  int B = out.size() / M / N;
+  bool use_matrix_kernels = has_global_scales && w_quantized &&
+      x.dtype() == bfloat16 && w_q.ndim() == 3 &&
+      K % (metal::is_nax_available() ? 64 : 32) == 0;
+
+  if (use_matrix_kernels && M == 1 && B >= 16 && right_sorted_) {
+    int E = w_q.size() / w_q.shape(-1) / w_q.shape(-2);
+    if (B / E >= 4) {
+      gather_qmm_rhs(
+          x,
+          w_q,
+          scales_w,
+          std::nullopt,
+          global_scale_w,
+          rhs_indices,
+          out,
+          true,
+          group_size_,
+          bits_,
+          x.size() / K,
+          N,
+          K,
+          d,
+          s,
+          mode);
+      return;
+    }
+  }
+
+  if (use_matrix_kernels && M >= get_qmv_batch_limit(K, N, d)) {
+    gather_qmm(
+        x,
+        w_q,
+        scales_w,
+        std::nullopt,
+        global_scale_w,
+        lhs_indices,
+        rhs_indices,
+        out,
+        true,
+        group_size_,
+        bits_,
+        M,
+        N,
+        K,
+        d,
+        s,
+        mode);
+    return;
+  }
+
   gather_qmv(
       x,
       w_q,

mlx/ops.cpp

@@ -5680,7 +5680,13 @@ array gather_qqmm(
   return array(
       std::move(out_shape),
       x.dtype(),
-      std::make_shared<GatherQQMM>(stream, group_size, bits, qmode),
+      std::make_shared<GatherQQMM>(
+          stream,
+          group_size,
+          bits,
+          qmode,
+          sorted_indices && !rhs_indices_,
+          sorted_indices && !lhs_indices_),
       std::move(inputs));
 }
 

python/tests/test_quantized.py

@@ -1555,6 +1555,74 @@ def test_gather_qqmm(self):
                 self.assertEqual(y_q.shape, y_hat.shape)
                 self.assertLess((y_q - y_hat).abs().max(), 1e-3)
 
+    @unittest.skipIf(mx.cuda.is_available(), "Not implemented for CUDA")
+    def test_gather_qqmm_global_scale_matrix_paths(self):
+        if mx.default_device() == mx.cpu:
+            self.skipTest("Not implemented for CPU")
+
+        E, N, K = 3, 64, 256
+        w = mx.random.normal((E, N, K), key=mx.random.key(10)).astype(mx.bfloat16)
+        global_scale_w = mx.max(mx.abs(w), axis=(1, 2)).astype(mx.float32)
+        quantized = [
+            mx.quantize(w[e], mode="nvfp4", global_scale=global_scale_w[e])
+            for e in range(E)
+        ]
+        w_q = mx.stack([q for q, _ in quantized])
+        scales_w = mx.stack([scales for _, scales in quantized])
+        w_hat = mx.stack(
+            [
+                mx.dequantize(
+                    q,
+                    scales,
+                    mode="nvfp4",
+                    dtype=mx.bfloat16,
+                    global_scale=global_scale_w[e],
+                )
+                for e, (q, scales) in enumerate(quantized)
+            ]
+        )
+
+        def check(x, rhs, lhs=None, sorted_indices=False):
+            global_scale_x = mx.max(mx.abs(x)).astype(mx.float32)
+            x_hat = mx.dequantize(
+                *mx.quantize(x, mode="nvfp4", global_scale=global_scale_x),
+                mode="nvfp4",
+                dtype=mx.bfloat16,
+                global_scale=global_scale_x,
+            )
+            actual = mx.gather_qqmm(
+                x,
+                w_q,
+                scales_w,
+                lhs,
+                rhs,
+                mode="nvfp4",
+                global_scale_x=global_scale_x,
+                global_scale_w=global_scale_w,
+                sorted_indices=sorted_indices,
+            )
+            expected = mx.gather_mm(
+                x_hat,
+                mx.swapaxes(w_hat, -1, -2),
+                lhs,
+                rhs,
+                sorted_indices=sorted_indices,
+            )
+            error = mx.abs(actual.astype(mx.float32) - expected.astype(mx.float32))
+            scale = mx.maximum(mx.abs(expected.astype(mx.float32)).max(), 1e-20)
+            self.assertLess(error.max() / scale, 3e-2)
+
+        # Sorted indices select the RHS-grouped matrix kernel.
+        x = mx.random.normal((24, 1, K), key=mx.random.key(20)).astype(mx.bfloat16)
+        rhs = mx.array([0] * 8 + [1] * 8 + [2] * 8)
+        check(x, rhs, sorted_indices=True)
+
+        # A large M selects the gathered matrix kernel.
+        x = mx.random.normal((2, 64, K), key=mx.random.key(30)).astype(mx.bfloat16)
+        lhs = mx.array([0, 1, 0, 1])
+        rhs = mx.array([0, 1, 2, 0])
+        check(x, rhs, lhs)
+
     def test_qmm_fp_type(self):
         indices = mx.array([[2], [0], [1]], dtype=mx.uint32)