From d20ba42575361f5efe067353e0282e5cd3a3e8ce Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Fri, 31 Jul 2026 17:13:50 -0700 Subject: [PATCH 1/3] Optimize large NVFP4 QMV on M5 Max Use a narrower output tile for large NVFP4 matrix-vector products on M5 Max GPUs. Reducing the results per SIMD-group from four to two lowers register pressure for output dimensions of 4096 or greater. Keep the existing kernel configuration for other architectures, quantization modes, and smaller output dimensions. Add coverage for batched and unbatched large-output QMV across the supported floating-point types. --- mlx/backend/metal/kernels/fp_quantized.h | 20 ++++++--- mlx/backend/metal/kernels/fp_quantized.metal | 41 ++++++++++++++---- mlx/backend/metal/quantized.cpp | 44 ++++++++++++++------ python/tests/test_quantized.py | 36 +++++++++++++++- 4 files changed, 114 insertions(+), 27 deletions(-) diff --git a/mlx/backend/metal/kernels/fp_quantized.h b/mlx/backend/metal/kernels/fp_quantized.h index f3aa24bf7a..6e77569f56 100644 --- a/mlx/backend/metal/kernels/fp_quantized.h +++ b/mlx/backend/metal/kernels/fp_quantized.h @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -321,7 +321,12 @@ METAL_FUNC void fp_qmv_quad_impl( } } -template +template < + typename T, + int group_size, + int bits, + bool has_global_scale = false, + int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( const device uint32_t* w, const device uint8_t* scales, @@ -335,7 +340,6 @@ METAL_FUNC void fp_qmv_fast_impl( uint simd_lid [[thread_index_in_simdgroup]]) { constexpr int packs_per_thread = 2; constexpr int num_simdgroups = 2; - constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack<32>(); constexpr int values_per_thread = pack_factor * packs_per_thread; @@ -1167,7 +1171,8 @@ template < int group_size, int bits, bool batched, - bool has_global_scale = false> + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( const device uint32_t* w, const device uint8_t* scales, @@ -1203,7 +1208,12 @@ template < s_strides, tid); } - fp_qmv_fast_impl( + fp_qmv_fast_impl< + T, + group_size, + bits, + has_global_scale, + results_per_simdgroup>( w, scales, global_scale, diff --git a/mlx/backend/metal/kernels/fp_quantized.metal b/mlx/backend/metal/kernels/fp_quantized.metal index 7404f2023f..d8d462288a 100644 --- a/mlx/backend/metal/kernels/fp_quantized.metal +++ b/mlx/backend/metal/kernels/fp_quantized.metal @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. // clang-format off #include "mlx/backend/metal/kernels/utils.h" @@ -57,6 +57,30 @@ aligned, \ batched) +#define instantiate_quantized_qmv_fast(mode, type, results, batched, group_size, bits) \ + instantiate_kernel( \ + #mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched, \ + fp_qmv_fast, \ + type, \ + group_size, \ + bits, \ + batched, \ + false, \ + results) \ + instantiate_kernel( \ + #mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched "_hgs", \ + fp_qmv_fast, \ + type, \ + group_size, \ + bits, \ + batched, \ + true, \ + results) + +#define instantiate_quantized_qmv_fast_r2(mode, type, group_size, bits) \ + instantiate_quantized_qmv_fast(mode, type, 2, 1, group_size, bits) \ + instantiate_quantized_qmv_fast(mode, type, 2, 0, group_size, bits) + #define instantiate_quantized_quad(mode, name, type, D, batched, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \ @@ -185,13 +209,14 @@ instantiate_quantized_all_rhs(type, mode, group_size, bits) #define instantiate_quantized_types(type) \ - instantiate_quantized_modes(type, nvfp4, 16, 4) \ - instantiate_quantized_modes(type, mxfp8, 32, 8) \ - instantiate_quantized_modes(type, mxfp4, 32, 4) \ - instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \ - instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \ - instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \ - instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \ + instantiate_quantized_modes(type, nvfp4, 16, 4) \ + instantiate_quantized_modes(type, mxfp8, 32, 8) \ + instantiate_quantized_modes(type, mxfp4, 32, 4) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \ + instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \ + instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \ + instantiate_quantized_qmv_fast_r2(nvfp4, type, 16, 4) instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 7c461bc1b5..7437237817 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include "mlx/backend/common/quantized.h" #include "mlx/backend/common/broadcasting.h" @@ -477,13 +477,19 @@ void qmv( int bn = 8; int bk = 32; - MTL::Size group_dims(bk, 2, 1); - MTL::Size grid_dims(M, (N + bn - 1) / bn, B); std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + // A narrower output tile reduces register pressure for large + // floating-point quantized matrix-vector products on M5 Max GPUs. + bool use_narrow_qmv = fast && N >= 4096 && d.get_architecture_gen() == 17 && + d.get_architecture().back() == 's' && mode == "nvfp4"; + int results_per_simdgroup = use_narrow_qmv ? 2 : 4; + bn = 2 * results_per_simdgroup; + MTL::Size group_dims(bk, 2, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); concatenate( kname, @@ -493,18 +499,30 @@ void qmv( group_size, "_b_", bits, + use_narrow_qmv ? "_r_2" : "", B > 1 ? "_batch_1" : "_batch_0", global_scale ? "_hgs" : ""); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - (fast ? "qmv_fast" : "qmv"), - mode, - type_string, - group_size, - bits, - B > 1, - global_scale.has_value()); + auto kernel = use_narrow_qmv ? get_quantized_kernel_wrapped( + d, + kname, + "qmv_fast", + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value(), + results_per_simdgroup) + : get_quantized_kernel_wrapped( + d, + kname, + (fast ? "qmv_fast" : "qmv"), + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value()); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index 0b756b39d5..feb97ab550 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -1,4 +1,4 @@ -# Copyright © 2023 Apple Inc. +# Copyright © 2023-2026 Apple Inc. import platform import subprocess @@ -516,6 +516,40 @@ def test_fp_qmv(self): self.assertEqual(y_q.shape, y_hat.shape) self.assertLess((y_q - y_hat).abs().max(), 1e-3) + def test_fp_qmv_large_output(self): + key = mx.random.key(0) + k1, k2 = mx.random.split(key) + K = 512 + N = 4096 + + for B in [1, 2]: + with self.subTest(B=B, N=N, K=K): + x_shape = (1, K) if B == 1 else (B, 1, K) + w_shape = (N, K) if B == 1 else (B, N, K) + x = mx.random.normal(shape=x_shape, key=k1) / K**0.5 + w = mx.random.normal(shape=w_shape, key=k2) + w_q, scales = mx.quantize(w, mode="nvfp4") + + dtypes = ( + [mx.float16, mx.bfloat16, mx.float32] + if mx.default_device() == mx.gpu + else [mx.float32] + ) + for dtype in dtypes: + with self.subTest(dtype=dtype): + x_t = x.astype(dtype) + w_hat = mx.dequantize(w_q, scales, mode="nvfp4", dtype=dtype) + y_q = mx.quantized_matmul( + x_t, + w_q, + scales, + transpose=True, + mode="nvfp4", + ) + y_hat = x_t @ mx.swapaxes(w_hat, -1, -2) + self.assertEqual(y_q.shape, y_hat.shape) + self.assertTrue(mx.allclose(y_q, y_hat, rtol=1e-3, atol=1e-3)) + def test_qmv_wide(self): # M in [2, vector_limit) routes to qmv_wide -- except K in {64, 128} # with power-of-2 bits, which stays on qmv_quad. Check both paths From 26386eb1ded08c35634debca91117c594476b839 Mon Sep 17 00:00:00 2001 From: Cheng Date: Sat, 8 Aug 2026 21:35:24 -0700 Subject: [PATCH 2/3] Simplify get_quantized_kernel_wrapped --- mlx/backend/metal/quantized.cpp | 32 +++++++++++--------------------- 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp index 7437237817..cb42e817d5 100644 --- a/mlx/backend/metal/quantized.cpp +++ b/mlx/backend/metal/quantized.cpp @@ -502,27 +502,17 @@ void qmv( use_narrow_qmv ? "_r_2" : "", B > 1 ? "_batch_1" : "_batch_0", global_scale ? "_hgs" : ""); - auto kernel = use_narrow_qmv ? get_quantized_kernel_wrapped( - d, - kname, - "qmv_fast", - mode, - type_string, - group_size, - bits, - B > 1, - global_scale.has_value(), - results_per_simdgroup) - : get_quantized_kernel_wrapped( - d, - kname, - (fast ? "qmv_fast" : "qmv"), - mode, - type_string, - group_size, - bits, - B > 1, - global_scale.has_value()); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + (fast ? "qmv_fast" : "qmv"), + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value(), + results_per_simdgroup); auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); From bc5d199c2b0e31e56d85dfa00a6afa509be04a14 Mon Sep 17 00:00:00 2001 From: Cheng Date: Sat, 8 Aug 2026 23:27:19 -0700 Subject: [PATCH 3/3] Fix cuda test failure --- python/tests/test_quantized.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tests/test_quantized.py b/python/tests/test_quantized.py index feb97ab550..806784cfc7 100644 --- a/python/tests/test_quantized.py +++ b/python/tests/test_quantized.py @@ -548,7 +548,8 @@ def test_fp_qmv_large_output(self): ) y_hat = x_t @ mx.swapaxes(w_hat, -1, -2) self.assertEqual(y_q.shape, y_hat.shape) - self.assertTrue(mx.allclose(y_q, y_hat, rtol=1e-3, atol=1e-3)) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(y_q, y_hat, rtol=tol, atol=tol)) def test_qmv_wide(self): # M in [2, vector_limit) routes to qmv_wide -- except K in {64, 128}