From af41c36e5e286c90a1440c127614069f99286342 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Fri, 4 Sep 2026 15:49:02 -0700 Subject: [PATCH 1/2] Use matrix kernels for global-scale gather_qqmm --- benchmarks/python/gather_qqmm_bench.py | 108 +++++++++++++++++++++++++ mlx/backend/metal/quantized.cpp | 51 ++++++++++++ mlx/ops.cpp | 8 +- python/tests/test_quantized.py | 67 +++++++++++++++ 4 files changed, 233 insertions(+), 1 deletion(-) create mode 100644 benchmarks/python/gather_qqmm_bench.py diff --git a/benchmarks/python/gather_qqmm_bench.py b/benchmarks/python/gather_qqmm_bench.py new file mode 100644 index 0000000000..83257a4cdb --- /dev/null +++ b/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) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index c8f8d9c443..baeda4ca8f 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -2148,6 +2148,57 @@ void GatherQQMM::eval_gpu(const std::vector& 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, diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 1a41688d28..7ec5369820 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -5680,7 +5680,13 @@ array gather_qqmm( return array( std::move(out_shape), x.dtype(), - std::make_shared(stream, group_size, bits, qmode), + std::make_shared( + stream, + group_size, + bits, + qmode, + sorted_indices && !rhs_indices_, + sorted_indices && !lhs_indices_), std::move(inputs)); } diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 051d81710c..76f67a053b 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1555,6 +1555,73 @@ def test_gather_qqmm(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + 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) From 8d19ced08d88882a00f1c76b9e705334f8eec729 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Fri, 11 Sep 2026 12:30:13 -0700 Subject: [PATCH 2/2] skip new test on CUDA --- python/tests/test_quantized.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 76f67a053b..0a52b6220e 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1555,6 +1555,7 @@ 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")