[Metal] Add affine qmv_mma for 3-16 rows

mlx/backend/metal/kernels/quantized.h

@@ -1916,6 +1916,187 @@ template <
       simd_lid);
 }
 
+template <typename T, int group_size, int bits, int WK, int BN, int TM>
+[[kernel]] void affine_qmv_mma(
+    const device uint32_t* w,
+    const device T* scales,
+    const device T* biases,
+    const device T* x,
+    device T* y,
+    const constant int& K,
+    const constant int& N,
+    const constant int& M,
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  static_assert(BN % 8 == 0, "BN should be divisible by 8");
+  static_assert(TM == 1 || TM == 2, "qmv_mma takes at most 16 rows");
+  static_assert(
+      (bits == 4 || bits == 8) &&
+          (group_size == 32 || group_size == 64 || group_size == 128),
+      "qmv_mma reads 4 and 8 bits at group sizes 32, 64 and 128");
+
+  using Frag = mlx::steel::BaseMMAFrag<float, 8, 8>;
+  using frag_t = typename Frag::frag_type;
+  constexpr int BK = 128;
+  constexpr short TN = BN / 8;
+  constexpr int pack_factor = 8 / bits;
+  constexpr short loads = bits / 4;
+
+  const short2 coord = Frag::get_coord(simd_lid);
+  const short fn = coord.x;
+  const short fm = coord.y;
+  const int row_bytes = K / pack_factor;
+  const int row_scales = K / group_size;
+
+  const int y_col = tid.x * BN;
+  auto wl = (const device uint8_t*)w +
+      (y_col + fm) * static_cast<int64_t>(row_bytes) + 16 * fn / pack_factor;
+  const int64_t s_off =
+      (y_col + fm) * static_cast<int64_t>(row_scales) + 16 * fn / group_size;
+  scales += s_off;
+  biases += s_off;
+  x += fn * static_cast<int64_t>(K) + 16 * fm;
+
+  bool x_rows[TM][2];
+  frag_t C[TN][TM];
+  STEEL_PRAGMA_UNROLL
+  for (short mt = 0; mt < TM; mt++) {
+    STEEL_PRAGMA_UNROLL
+    for (short j = 0; j < 2; j++) {
+      x_rows[mt][j] = 8 * mt + fn + j < M;
+    }
+    STEEL_PRAGMA_UNROLL
+    for (short t = 0; t < TN; t++) {
+      C[t][mt] = frag_t(0);
+    }
+  }
+
+  for (int k = simd_gid * BK; k < K; k += WK * BK) {
+    vec<T, 8> xs[TM][2][2];
+    STEEL_PRAGMA_UNROLL
+    for (short mt = 0; mt < TM; mt++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < 2; j++) {
+        const device T* xr = x + (8 * mt + j) * static_cast<int64_t>(K) + k;
+        STEEL_PRAGMA_UNROLL
+        for (short h = 0; h < 2; h++) {
+          xs[mt][j][h] = x_rows[mt][j] ? *(const device vec<T, 8>*)(xr + 8 * h)
+                                       : vec<T, 8>(0);
+        }
+      }
+    }
+    uint4 ws[TN][loads];
+    float sc[TN];
+    float bs[TN];
+    STEEL_PRAGMA_UNROLL
+    for (short t = 0; t < TN; t++) {
+      const device uint8_t* wt =
+          wl + 8 * t * static_cast<int64_t>(row_bytes) + k / pack_factor;
+      STEEL_PRAGMA_UNROLL
+      for (short l = 0; l < loads; l++) {
+        ws[t][l] = *(const device uint4*)(wt + 16 * l);
+      }
+      const int64_t si =
+          8 * t * static_cast<int64_t>(row_scales) + k / group_size;
+      sc[t] = float(scales[si]);
+      bs[t] = float(biases[si]);
+    }
+
+    STEEL_PRAGMA_UNROLL
+    for (short g = 0; g < 4; g++) {
+      frag_t b[4][TM];
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < 4; i++) {
+        const short f = 4 * g + i;
+        STEEL_PRAGMA_UNROLL
+        for (short mt = 0; mt < TM; mt++) {
+          b[i][mt] = frag_t(
+              float(xs[mt][0][f / 8][f % 8]), float(xs[mt][1][f / 8][f % 8]));
+        }
+      }
+      STEEL_PRAGMA_UNROLL
+      for (short t = 0; t < TN; t++) {
+        // p[i] holds weights 4g+i and 16+4g+i of the lane. With 0x6400 set
+        // each half reads as 1024 + q, and taking 1024 off again is exact.
+        uint p[4];
+        if constexpr (bits == 4) {
+          const uint lo = ws[t][0][g / 2];
+          const uint hi = ws[t][0][2 + g / 2];
+          const uint q = (g % 2) == 0 ? (lo & 0xffff) | (hi << 16)
+                                      : (lo >> 16) | (hi & 0xffff0000);
+          STEEL_PRAGMA_UNROLL
+          for (short i = 0; i < 4; i++) {
+            p[i] = (q >> (4 * i)) & 0x000f000f;
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short i = 0; i < 4; i++) {
+            p[i] = ((ws[t][0][g] >> (8 * i)) & 0xff) |
+                (((ws[t][1][g] >> (8 * i)) & 0xff) << 16);
+          }
+        }
+        STEEL_PRAGMA_UNROLL
+        for (short i = 0; i < 4; i++) {
+          const half2 v = as_type<half2>(p[i] | 0x64006400) - half2(1024.0h);
+          frag_t a(
+              fma(float(v.x), sc[t], bs[t]), fma(float(v.y), sc[t], bs[t]));
+          STEEL_PRAGMA_UNROLL
+          for (short mt = 0; mt < TM; mt++) {
+            Frag::mma(C[t][mt], a, b[i][mt], C[t][mt]);
+          }
+        }
+      }
+    }
+  }
+
+  constexpr int tile_elems = TN * TM * 2;
+  threadgroup float Ps[(WK - 1) * tile_elems * 32];
+  if (simd_gid > 0) {
+    threadgroup float* p = Ps + (simd_gid - 1) * tile_elems * 32 + simd_lid;
+    STEEL_PRAGMA_UNROLL
+    for (short t = 0; t < TN; t++) {
+      STEEL_PRAGMA_UNROLL
+      for (short mt = 0; mt < TM; mt++) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < 2; j++) {
+          p[((t * TM + mt) * 2 + j) * 32] = C[t][mt][j];
+        }
+      }
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd_gid == 0) {
+    for (int s = 0; s < WK - 1; s++) {
+      const threadgroup float* p = Ps + s * tile_elems * 32 + simd_lid;
+      STEEL_PRAGMA_UNROLL
+      for (short t = 0; t < TN; t++) {
+        STEEL_PRAGMA_UNROLL
+        for (short mt = 0; mt < TM; mt++) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < 2; j++) {
+            C[t][mt][j] += p[((t * TM + mt) * 2 + j) * 32];
+          }
+        }
+      }
+    }
+    STEEL_PRAGMA_UNROLL
+    for (short mt = 0; mt < TM; mt++) {
+      STEEL_PRAGMA_UNROLL
+      for (short j = 0; j < 2; j++) {
+        if (x_rows[mt][j]) {
+          device T* yr =
+              y + (8 * mt + fn + j) * static_cast<int64_t>(N) + y_col + fm;
+          STEEL_PRAGMA_UNROLL
+          for (short t = 0; t < TN; t++) {
+            yr[8 * t] = static_cast<T>(C[t][mt][j]);
+          }
+        }
+      }
+    }
+  }
+}
+
 template <
     typename T,
     int group_size,

mlx/backend/metal/kernels/quantized.metal

@@ -195,6 +195,26 @@
   instantiate_qmv_wide_batched(affine_qmv_wide, type, group_size, bits, 7, 8) \
   instantiate_qmv_wide_batched(affine_qmv_wide, type, group_size, bits, 8, 8)
 
+///////////////////////////////////////////////////////////////////////////////
+// affine_qmv_mma
+///////////////////////////////////////////////////////////////////////////////
+
+#define instantiate_qmv_mma(name, type, group_size, bits, wk, bn, tm) \
+  instantiate_kernel( \
+      #name "_" #type "_gs_" #group_size "_b_" #bits "_wk_" #wk "_bn_" #bn "_tm_" #tm, \
+      name,           \
+      type,           \
+      group_size,     \
+      bits,           \
+      wk,             \
+      bn,             \
+      tm)
+
+#define instantiate_qmv_mma_all(type, group_size, bits)               \
+  instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 16, 1) \
+  instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 16, 2) \
+  instantiate_qmv_mma(affine_qmv_mma, type, group_size, bits, 4, 32, 1)
+
 ///////////////////////////////////////////////////////////////////////////////
 // Instantiation
 ///////////////////////////////////////////////////////////////////////////////
@@ -216,7 +236,9 @@
   instantiate_quantized_all(type, group_size, 4)         \
   instantiate_quantized_all(type, group_size, 5)         \
   instantiate_quantized_all(type, group_size, 6)         \
-  instantiate_quantized_all(type, group_size, 8)
+  instantiate_quantized_all(type, group_size, 8)         \
+  instantiate_qmv_mma_all(type, group_size, 4)           \
+  instantiate_qmv_mma_all(type, group_size, 8)
 
 #define instantiate_quantized_all_groups(type) \
   instantiate_quantized_all_bits(type, 32)     \

mlx/backend/metal/quantized.cpp

@@ -641,7 +641,15 @@ inline bool use_qmv_mma(
     int N,
     int K,
     metal::Device& d) {
-  if (K % 128 != 0) {
+  bool nax = metal::is_nax_available();
+  bool affine = mode == QuantizationMode::Affine;
+  // affine_qmv_mma beats qmv from 3 rows below gen 15 and qmv_wide from 4
+  // rows above; the other paths from 6.
+  int min_rows = 6;
+  if (affine && !nax) {
+    min_rows = d.get_architecture_gen() >= 15 ? 4 : 3;
+  }
+  if ((M < min_rows) || (K % 128 != 0)) {
     return false;
   }
   // Require aligned vector load.
@@ -650,12 +658,22 @@ inline bool use_qmv_mma(
     return false;
   }
   // Route for large K/N.
-  bool nax = metal::is_nax_available();
-  if (mode == QuantizationMode::Affine) {
-    // Only qmv_nax reads affine weights, with 3 to 8 bits.
-    return nax && (bits >= 3 && bits != 7) && (group_size >= 32) &&
-        (x.dtype() == float16 || x.dtype() == bfloat16) && (N >= 512) &&
-        (K >= 1024) && (N % 32 == 0);
+  if (affine) {
+    if (nax) {
+      // qmv_nax reads affine weights with 3 to 8 bits.
+      return (bits >= 3 && bits != 7) &&
+          (x.dtype() == float16 || x.dtype() == bfloat16) && (N >= 512) &&
+          (K >= 1024) && (N % 32 == 0);
+    } else { // non-nax
+      // Ultra keeps the vector kernels.
+      if (d.get_architecture().back() == 'd') {
+        return false;
+      }
+      // Only 4 and 8 bits are instantiated; affine pays off from a smaller N *
+      // K than fp.
+      return (bits == 4 || bits == 8) && (N >= 512) && (K >= 512) &&
+          (int64_t(N) * K >= (int64_t(1) << 21)) && (N % 16 == 0);
+    }
   } else { // fp quantization
     if (nax) {
       // Below 2048, nvfp4 with few rows is faster on qmv_wide, which runs for
@@ -690,8 +708,25 @@ void qmv_mma(
     const Stream& s,
     const std::string& mode) {
   bool nax = metal::is_nax_available();
-  int wk = nax || d.get_architecture_gen() >= 15 ? 8 : 4;
-  int bn = nax ? 32 : 16;
+  bool affine = (mode == "affine");
+  int wk = 4;
+  int bn = 16;
+  int tm = (M > 8) ? 2 : 1;
+  if (nax) {
+    wk = 8;
+    bn = 32;
+  } else { // non-nax
+    // Affine stays at wk 4 on gen-15+ too (wk 8 was no faster on M4).
+    if (!affine && (d.get_architecture_gen() >= 15)) {
+      wk = 8;
+    }
+    // Below gen 15, affine with one row tile is faster at 32 columns (M2 Pro).
+    if (affine && (d.get_architecture_gen() < 15) && (tm == 1) &&
+        (N % 32 == 0)) {
+      bn = 32;
+    }
+  }
+
   std::string type_string = get_type_string(x.dtype());
   std::string kname;
   kname.reserve(64);
@@ -713,7 +748,6 @@ void qmv_mma(
     kernel = get_qmm_nax_kernel_wrapped(
         d, kname, "qmv_nax", mode, type_string, group_size, bits, wk, bn);
   } else {
-    int tm = M > 8 ? 2 : 1;
     concatenate(kname, "_tm_", tm);
     kernel = get_quantized_kernel_wrapped(
         d, kname, "qmv_mma", mode, type_string, group_size, bits, wk, bn, tm);
@@ -2184,7 +2218,7 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto mode = quantization_mode_to_string(mode_);
 
   // For few rows use qmv_mma when N/K is large enough.
-  if (transpose_ && non_batched && (M >= 6) && (M <= 16) &&
+  if (transpose_ && non_batched && (M <= 16) &&
       use_qmv_mma(x, w, scales, mode_, group_size_, bits_, M, N, K, d)) {
     qmv_mma(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode);
     return;

python/tests/test_quantized.py

@@ -849,7 +849,7 @@ def rel_err(got, expected):
                 self.assertLess(rel_err(out, x @ w_hat.swapaxes(-1, -2)), 3e-2)
 
     @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
-    def test_qmm_few_rows(self):
+    def test_fp_qmm_few_rows(self):
         if mx.default_device() == mx.cpu:
             self.skipTest("Covers GPU kernels only")
         mx.random.seed(0)
@@ -884,6 +884,43 @@ def rel_err(got, expected):
                     out = mx.quantized_matmul(xv, w_q, sv, mode=mode)
                     self.assertLess(rel_err(out, xv @ w_hat.T), tol)
 
+    @unittest.skipIf(not mx.metal.is_available(), "requires Metal")
+    def test_affine_qmm_few_rows(self):
+        if mx.default_device() == mx.cpu:
+            self.skipTest("Covers GPU kernels only")
+        mx.random.seed(0)
+
+        def rel_err(got, expected):
+            d = mx.abs(got.astype(mx.float32) - expected.astype(mx.float32))
+            return d.max() / mx.abs(expected.astype(mx.float32)).max()
+
+        for (K, N), M, bits, gs, dtype in product(
+            [
+                (4224, 4128),
+                (512, 4096),  # smallest K: one K step per simdgroup
+            ],
+            [3, 8, 9],
+            [4, 8],
+            [32, 64, 128],
+            [mx.bfloat16, mx.float16, mx.float32],
+        ):
+            with self.subTest(M=M, N=N, K=K, bits=bits, gs=gs, dtype=dtype):
+                tol = 1e-5 if dtype == mx.float32 else 3e-2
+                x = (mx.random.normal((M, K)) / K**0.5).astype(dtype)
+                w = mx.random.normal((N, K), dtype=dtype)
+                w_q, scales, biases = mx.quantize(w, group_size=gs, bits=bits)
+                w_hat = mx.dequantize(w_q, scales, biases, group_size=gs, bits=bits)
+                y_ref = x @ w_hat.T
+                out = mx.quantized_matmul(
+                    x,
+                    w_q,
+                    scales,
+                    biases,
+                    group_size=gs,
+                    bits=bits,
+                )
+                self.assertLess(rel_err(out, y_ref), tol)
+
     def test_qmm_vjp(self):
         key = mx.random.key(0)
         k1, k2 = mx.random.split(key)