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..cb42e817d5 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,6 +499,7 @@ 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( @@ -504,7 +511,8 @@ void qmv( group_size, bits, B > 1, - global_scale.has_value()); + global_scale.has_value(), + results_per_simdgroup); 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..806784cfc7 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,41 @@ 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) + 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} # with power-of-2 bits, which stays on qmv_quad. Check both paths