Optimize large NVFP4 QMV on M5 Max

mlx/backend/metal/kernels/fp_quantized.h

@@ -1,4 +1,4 @@
-// Copyright © 2025 Apple Inc.
+// Copyright © 2025-2026 Apple Inc.
 
 #include <metal_simdgroup>
 #include <metal_stdlib>
@@ -321,7 +321,12 @@ METAL_FUNC void fp_qmv_quad_impl(
   }
 }
 
-template <typename T, int group_size, int bits, bool has_global_scale = false>
+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<T, group_size, bits, has_global_scale>(
+  fp_qmv_fast_impl<
+      T,
+      group_size,
+      bits,
+      has_global_scale,
+      results_per_simdgroup>(
       w,
       scales,
       global_scale,

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)

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);

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