Add Metal SDPA support for D96/V64

mlx/backend/metal/jit_kernels.cpp

@@ -1360,6 +1360,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel(
     int bq,
     int bk,
     int bd,
+    int bv,
     int wm,
     int wn,
     const array& m,
@@ -1371,16 +1372,28 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel(
         kernel_source,
         metal::utils(),
         metal::steel_attention_nax(),
-        get_template_definition(
-            lib_name,
-            split_d ? "attention_nax_dsplit" : "attention_nax",
-            get_type_string(q.dtype()),
-            bq,
-            bk,
-            bd,
-            wm,
-            wn,
-            get_type_string(m.dtype())));
+        split_d ? get_template_definition(
+                      lib_name,
+                      "attention_nax_dsplit",
+                      get_type_string(q.dtype()),
+                      bq,
+                      bk,
+                      bd,
+                      wm,
+                      wn,
+                      get_type_string(m.dtype()))
+                : get_template_definition(
+                      lib_name,
+                      "attention_nax",
+                      get_type_string(q.dtype()),
+                      bq,
+                      bk,
+                      bd,
+                      wm,
+                      wn,
+                      get_type_string(m.dtype()),
+                      "float",
+                      bv));
     return kernel_source;
   });
   return d.get_kernel(kernel_name, lib, hash_name, func_consts);

mlx/backend/metal/kernels.h

@@ -426,6 +426,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel(
     int bq,
     int bk,
     int bd,
+    int bv,
     int wm,
     int wn,
     const array& m,

mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -40,6 +40,7 @@ using namespace metal;
 
 #define instantiate_sdpa_vector_heads(type)      \
   instantiate_sdpa_vector(type, 64, 64)          \
+  instantiate_sdpa_vector(type, 96, 64)          \
   instantiate_sdpa_vector(type, 96, 96)          \
   instantiate_sdpa_vector(type, 128, 128)        \
   instantiate_sdpa_vector(type, 192, 128)        \

mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h

@@ -79,7 +79,8 @@ template <
     int WM,
     int WN,
     typename MaskType = float,
-    typename AccumType = float>
+    typename AccumType = float,
+    int BV = BD>
 [[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention_nax(
     const device T* Q [[buffer(0)]],
     const device T* K [[buffer(1)]],
@@ -136,11 +137,13 @@ template <
   constexpr int TQ = BQ / (kNWarps * kU);
   // HeadDim frags (all warps load the same frags)
   constexpr int TD = BD / kU;
+  constexpr int TV = BV / kU;
+  static_assert(BD % kU == 0 && BV % (2 * kU) == 0, "Invalid head dimensions");
   // KV seq frags per warp
   constexpr short TK = BK / kU;
 
   static_assert(TQ == 1, "Check TQ");
-  using otile_t = NAXTile<AccumType, TQ, TD>;
+  using otile_t = NAXTile<AccumType, TQ, TV>;
   otile_t Otile;
 
   Otile.clear();
@@ -421,8 +424,8 @@ template <
     STEEL_PRAGMA_UNROLL
     for (short iq = 0; iq < TQ; iq++) {
       STEEL_PRAGMA_UNROLL
-      for (short id = 0; id < TD; id += 2) {
-        if constexpr (BD == 128) {
+      for (short id = 0; id < TV; id += 2) {
+        if constexpr (BV == 128) {
           if (id == 4) {
             threadgroup_barrier(mem_flags::mem_none);
           }

mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal

@@ -11,6 +11,12 @@
       "_wm" #wm "_wn" #wn "_mask" #mname,                                \
   attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float)
 
+#define instantiate_attn_asymmetric(tname, dtype, bq, bk, bd, bv, wm, wn, mname, mtype) \
+  instantiate_kernel(                                                               \
+      "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd "_bv" #bv             \
+      "_wm" #wm "_wn" #wn "_mask" #mname,                                            \
+  attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float, bv)
+
 #define instantiate_attn_dsplit(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \
   instantiate_kernel(                                                           \
       "steel_attention_dsplit_" #tname "_bq" #bq "_bk" #bk "_bd" #bd            \
@@ -27,7 +33,9 @@
 
 #define instantiate_attn_mask_helper(iname, itype) \
     instantiate_attn_shapes_helper(iname, itype, iname, itype) \
-    instantiate_attn_shapes_helper(iname, itype, bool_, bool)
+    instantiate_attn_shapes_helper(iname, itype, bool_, bool) \
+    instantiate_attn_asymmetric(iname, itype, 64, 32, 96, 64, 4, 1, iname, itype) \
+    instantiate_attn_asymmetric(iname, itype, 64, 32, 96, 64, 4, 1, bool_, bool)
 
 instantiate_attn_mask_helper(float16, half);
 instantiate_attn_mask_helper(bfloat16, bfloat);

mlx/backend/metal/nojit_kernels.cpp

@@ -505,6 +505,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel(
     int,
     int,
     int,
+    int,
     const array&,
     bool) {
   return d.get_kernel(kernel_name, hash_name, func_consts);

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -29,6 +29,7 @@ void sdpa_full_self_attention_nax(
   using namespace mlx::steel;
 
   int bd = q.shape(-1);
+  int bv = v.shape(-1);
   int bq = 64;
   int bk = 32;
 
@@ -98,6 +99,7 @@ void sdpa_full_self_attention_nax(
       bk,
       "_bd",
       bd,
+      (bv == bd ? "" : "_bv" + std::to_string(bv)),
       "_wm",
       wm,
       "_wn",
@@ -131,6 +133,7 @@ void sdpa_full_self_attention_nax(
       bq,
       bk,
       bd,
+      bv,
       wm,
       wn,
       (has_mask ? *mask : q),
@@ -733,17 +736,30 @@ std::tuple<bool, std::string> has_fused_kernel(
 
   std::ostringstream msg;
   if (query_sequence_length > 8) {
-    const bool supported_head_dim = query_head_dim == value_head_dim &&
-        (query_head_dim == 64 || query_head_dim == 72 || query_head_dim == 80 ||
-         query_head_dim == 96 || query_head_dim == 128 ||
-         query_head_dim == 192 || query_head_dim == 256);
+    const bool asymmetric = query_head_dim == 96 && value_head_dim == 64;
+    const bool supported_head_dim = asymmetric ||
+        (query_head_dim == value_head_dim &&
+         (query_head_dim == 64 || query_head_dim == 72 ||
+          query_head_dim == 80 || query_head_dim == 96 ||
+          query_head_dim == 128 || query_head_dim == 192 ||
+          query_head_dim == 256));
     if (!supported_head_dim) {
       msg << "the full attention kernel supports head dims "
           << "{64, 72, 80, 96, 128, 192, 256} with matching query/value head "
-          << "dims; got query head dim " << query_head_dim
-          << " and value head dim " << value_head_dim << ".";
+          << "dims, or (query, value) head dims (96, 64) with NAX; got query "
+          << "head dim " << query_head_dim << " and value head dim "
+          << value_head_dim << ".";
       return {false, msg.str()};
     }
+    if (asymmetric) {
+      if (!metal::is_nax_available() ||
+          (q.dtype() == float32 && !env::enable_tf32())) {
+        return {
+            false,
+            "the (96, 64) full attention kernel requires NAX and "
+            "TF32 for float32 inputs."};
+      }
+    }
     if (has_mask && !has_arr_mask &&
         !(query_sequence_length <= key_sequence_length && do_causal)) {
       msg << "the full attention kernel with a causal mask requires the "
@@ -758,11 +774,12 @@ std::tuple<bool, std::string> has_fused_kernel(
          (query_head_dim == 64 || query_head_dim == 96 ||
           query_head_dim == 128 || query_head_dim == 192 ||
           query_head_dim == 256 || query_head_dim == 512)) ||
-        (query_head_dim == 192 && value_head_dim == 128);
+        (query_head_dim == 192 && value_head_dim == 128) ||
+        (query_head_dim == 96 && value_head_dim == 64);
     if (!supported_head_dim) {
       msg << "the vector attention kernel supports head dims "
           << "{64, 96, 128, 192, 256, 512} with matching query/value head "
-          << "dims, or query head dim 192 with value head dim 128; got "
+          << "dims, or (query, value) head dims (96, 64) or (192, 128); got "
           << "query head dim " << query_head_dim << " and value head dim "
           << value_head_dim << ".";
       return {false, msg.str()};

python/tests/test_fast_sdpa.py

@@ -1,5 +1,6 @@
 import math
 import os
+import platform
 import unittest
 from itertools import product
 
@@ -292,6 +293,92 @@ def test_sdpa_full_head_dim_256(self):
                         tol = 5e-3
                     self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol))
 
+    @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
+    def test_sdpa_full_d96_v64(self):
+        if mx.default_device() != mx.gpu:
+            self.skipTest("requires GPU")
+        arch = mx.device_info()["architecture"]
+        min_gen = 18 if arch.endswith("p") else 17
+        if int(arch[-3:-1]) < min_gen or tuple(
+            map(int, platform.mac_ver()[0].split(".")[:2])
+        ) < (26, 2):
+            self.skipTest("requires NAX")
+        mx.random.seed(0)
+        scale = 96**-0.5
+        cells = [
+            (1, 40, 40, 9, 33),
+            (1, 4, 4, 64, 64),
+            (2, 4, 2, 65, 97),
+            (1, 4, 4, 16, 9),
+        ]
+        for dtype, (B, qH, kH, qL, kL) in product(
+            (mx.float16, mx.bfloat16, mx.float32), cells
+        ):
+            if dtype == mx.float32 and os.environ.get("MLX_ENABLE_TF32", "1") == "0":
+                continue
+            q = (
+                (0.5 * mx.random.normal((B, qL, qH, 96)))
+                .astype(dtype)
+                .transpose(0, 2, 1, 3)
+            )
+            k = (0.5 * mx.random.normal((B, kH, kL + 32, 96))).astype(dtype)[:, :, :kL]
+            masks = [
+                None,
+                mx.arange(kL) < kL - 3,
+                mx.random.normal((qL, kL)).astype(dtype),
+            ]
+            if qL <= kL:
+                masks.append("causal")
+            for dv, mask, with_sinks in product((64, 96), masks, (False, True)):
+                with self.subTest(
+                    dtype=dtype,
+                    cell=(B, qH, kH, qL, kL),
+                    dv=dv,
+                    mask=type(mask).__name__,
+                    sinks=with_sinks,
+                ):
+                    v = (0.5 * mx.random.normal((B, kH, kL + 32, dv))).astype(dtype)[
+                        :, :, :kL
+                    ]
+                    if qL == 9:
+                        # V sliced from the interleaved K/V projection.
+                        v = (0.5 * mx.random.normal((B, kL, kH, 2 * dv))).astype(dtype)
+                        v = v.transpose(0, 2, 1, 3)[..., dv:]
+                    sinks = (
+                        mx.random.normal((qH,)).astype(dtype) if with_sinks else None
+                    )
+                    ref = mlx_ref_attn(
+                        q.astype(mx.float32),
+                        k.astype(mx.float32),
+                        v.astype(mx.float32),
+                        scale,
+                        mask,
+                        sinks,
+                    )
+                    out = mx.fast.scaled_dot_product_attention(
+                        q,
+                        k,
+                        v,
+                        scale=scale,
+                        mask=mask,
+                        sinks=sinks,
+                        force_fused=True,
+                    )
+                    self.assertEqual(out.shape, (B, qH, qL, dv))
+                    tol = 5e-3 if dtype == mx.bfloat16 else 1e-3
+                    self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol))
+
+    @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
+    def test_sdpa_full_d96_rejects_unsupported_value_dim(self):
+        if mx.default_device() != mx.gpu:
+            self.skipTest("requires GPU")
+        q = mx.random.normal((1, 4, 9, 96)).astype(mx.float16)
+        k = mx.random.normal((1, 4, 33, 96)).astype(mx.float16)
+        v = mx.random.normal((1, 4, 33, 80)).astype(mx.float16)
+        scale = 96**-0.5
+        with self.assertRaisesRegex(ValueError, "supports head dims"):
+            mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, force_fused=True)
+
     def test_sdpa_vector_kv_transposed_head_seq(self):
         D = 64
         Nq = 4
@@ -410,6 +497,51 @@ def test_sdpa_vector_gqa_long(self):
                 out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
                 self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
 
+    @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
+    def test_sdpa_vector_d96_v64(self):
+        if mx.default_device() != mx.gpu:
+            self.skipTest("requires GPU")
+        mx.random.seed(0)
+        scale = 96**-0.5
+        # Include MiniCPM3 heads and GQA to reach two-pass on smaller GPUs.
+        cells = [(1, 40, 40, 1, 127), (1, 40, 40, 8, 1025), (2, 4, 2, 4, 4097)]
+        for dtype, (B, qH, kH, qL, kL) in product(
+            (mx.float32, mx.float16, mx.bfloat16), cells
+        ):
+            q = mx.random.normal((B, qL, qH, 96)).astype(dtype)
+            q = q.transpose(0, 2, 1, 3)
+            k = mx.random.normal((B, kH, kL + 32, 96)).astype(dtype)[:, :, :kL]
+            v = mx.random.normal((B, kH, kL + 32, 64)).astype(dtype)[:, :, :kL]
+            masks = (
+                None,
+                "causal",
+                mx.arange(kL) < kL - 3,
+                mx.random.normal((qL, kL)).astype(dtype),
+            )
+            for mask in masks:
+                with self.subTest(
+                    dtype=dtype, cell=(B, qH, kH, qL, kL), mask=type(mask).__name__
+                ):
+                    # GPU float32 GEMM may use TF32; keep the reference exact.
+                    with mx.stream(mx.cpu):
+                        ref = mlx_ref_attn(
+                            q.astype(mx.float32),
+                            k.astype(mx.float32),
+                            v.astype(mx.float32),
+                            scale,
+                            mask,
+                        )
+                    out = mx.fast.scaled_dot_product_attention(
+                        q, k, v, scale=scale, mask=mask, force_fused=True
+                    )
+                    self.assertEqual(out.shape, (B, qH, qL, 64))
+                    tol = (
+                        1e-5
+                        if dtype == mx.float32
+                        else (3e-3 if dtype == mx.float16 else 2e-2)
+                    )
+                    self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol))
+
     @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
     def test_sdpa_vector_head_dim_512(self):
         if mx.default_device() != mx.gpu: