Add D512 support to Metal vector attention

mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -45,6 +45,7 @@ using namespace metal;
   instantiate_sdpa_vector(type, 192, 128)        \
   instantiate_sdpa_vector(type, 192, 192)        \
   instantiate_sdpa_vector(type, 256, 256)        \
+  instantiate_sdpa_vector(type, 512, 512)        \
   instantiate_sdpa_vector_gqa(type, 64, 64, 8, 8)     \
   instantiate_sdpa_vector_gqa(type, 128, 128, 8, 4)   \
   instantiate_sdpa_vector_gqa(type, 128, 128, 12, 4)  \
@@ -53,7 +54,8 @@ using namespace metal;
   instantiate_sdpa_vector_aggregation(type, 96)  \
   instantiate_sdpa_vector_aggregation(type, 128) \
   instantiate_sdpa_vector_aggregation(type, 192) \
-  instantiate_sdpa_vector_aggregation(type, 256)
+  instantiate_sdpa_vector_aggregation(type, 256) \
+  instantiate_sdpa_vector_aggregation(type, 512)
 
 instantiate_sdpa_vector_heads(float)
 instantiate_sdpa_vector_heads(bfloat16_t)

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -757,13 +757,13 @@ std::tuple<bool, std::string> has_fused_kernel(
         (query_head_dim == value_head_dim &&
          (query_head_dim == 64 || query_head_dim == 96 ||
           query_head_dim == 128 || query_head_dim == 192 ||
-          query_head_dim == 256)) ||
+          query_head_dim == 256 || query_head_dim == 512)) ||
         (query_head_dim == 192 && value_head_dim == 128);
     if (!supported_head_dim) {
       msg << "the vector attention kernel supports head dims "
-          << "{64, 96, 128, 192, 256} with matching query/value head dims, "
-          << "or query head dim 192 with value head dim 128; got query head "
-          << "dim " << query_head_dim << " and value head dim "
+          << "{64, 96, 128, 192, 256, 512} with matching query/value head "
+          << "dims, or query head dim 192 with value head dim 128; got "
+          << "query head dim " << query_head_dim << " and value head dim "
           << value_head_dim << ".";
       return {false, msg.str()};
     }
@@ -780,6 +780,36 @@ std::tuple<bool, std::string> has_fused_kernel(
           << query_sequence_length << " and GQA factor " << gqa_factor << ".";
       return {false, msg.str()};
     }
+    // Head dim 512 uses the generic vector kernel, reading K/V is its main
+    // cost. By default, use it only for one query, GQA factor 8, and no array
+    // mask.
+    if (query_head_dim == 512) {
+      int min_key_sequence_length = env::get_var("MLX_SDPA_D512_MIN_KL", 1024);
+      bool always = (min_key_sequence_length == 0);
+      if (!always && query_sequence_length != 1) {
+        msg << "the vector attention kernel for head dim 512 defaults to "
+               "single-token queries; got query length "
+            << query_sequence_length << ".";
+        return {false, msg.str()};
+      }
+      if (!always && gqa_factor != 8) {
+        msg << "the vector attention kernel for head dim 512 defaults to "
+               "GQA factor 8; got GQA factor "
+            << gqa_factor << ".";
+        return {false, msg.str()};
+      }
+      if (!always && has_arr_mask) {
+        msg << "the vector attention kernel for head dim 512 defaults to "
+               "calls without array masks.";
+        return {false, msg.str()};
+      }
+      if (key_sequence_length < min_key_sequence_length) {
+        msg << "the vector attention kernel for head dim 512 requires at "
+            << "least " << min_key_sequence_length << " keys; got key length "
+            << key_sequence_length << ".";
+        return {false, msg.str()};
+      }
+    }
   }
   return {true, ""};
 }

python/tests/test_fast_sdpa.py

@@ -1,7 +1,9 @@
 import math
 import os
 import unittest
+from contextlib import nullcontext
 from itertools import product
+from unittest.mock import patch
 
 import mlx.core as mx
 import mlx_tests
@@ -418,6 +420,85 @@ 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_head_dim_512(self):
+        if mx.default_device() != mx.gpu:
+            self.skipTest("requires GPU")
+        # gemma-4 global attention: 32 query heads over 4 key/value heads.
+        D = 512
+        Nq, Nkv = 32, 4
+        scale = D**-0.5
+        mx.random.seed(0)
+
+        # Keep the test hermetic if the caller already configured a threshold.
+        # addCleanup restores the complete environment snapshot on every exit.
+        environment = patch.dict(os.environ)
+        environment.start()
+        self.addCleanup(environment.stop)
+        os.environ["MLX_SDPA_D512_MIN_KL"] = "0"
+
+        # Test 1-pass kernel.
+        for dtype in (mx.float32, mx.float16, mx.bfloat16):
+            with self.subTest(L=128, dtype=dtype, threshold="off"):
+                q = mx.random.normal(shape=(1, Nq, 1, D), dtype=dtype)
+                k = mx.random.normal(shape=(1, Nkv, 128, D), dtype=dtype)
+                v = mx.random.normal(shape=(1, Nkv, 128, D), dtype=dtype)
+                ref = mlx_ref_attn(q, k, v, scale)
+                out = mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=scale, force_fused=True
+                )
+                atol = 1e-5 if dtype == mx.float32 else 2e-2
+                self.assertTrue(mx.allclose(ref, out, atol=atol))
+
+        # Test 2-pass kernel.
+        for L, dtype in product((8192, 8201), (mx.float32, mx.float16, mx.bfloat16)):
+            with self.subTest(L=L, dtype=dtype):
+                q = mx.random.normal(shape=(1, Nq, 1, D), dtype=dtype)
+                k = mx.random.normal(shape=(1, Nkv, L, D), dtype=dtype)
+                v = mx.random.normal(shape=(1, Nkv, L, D), dtype=dtype)
+                ref = mlx_ref_attn(q, k, v, scale)
+                out = mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=scale, force_fused=True
+                )
+                atol = 1e-5 if dtype == mx.float32 else 2e-2
+                self.assertTrue(mx.allclose(ref, out, atol=atol))
+
+        # Test other heads.
+        for q_heads, kv_heads in ((8, 8), (32, 8)):
+            with self.subTest(q_heads=q_heads, kv_heads=kv_heads):
+                q = mx.random.normal(shape=(1, q_heads, 1, D))
+                k = mx.random.normal(shape=(1, kv_heads, L, D))
+                v = mx.random.normal(shape=(1, kv_heads, L, D))
+                ref = mlx_ref_attn(q, k, v, scale)
+                out = mx.fast.scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    scale=scale,
+                    force_fused=True,
+                )
+                self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
+
+        # Test batched.
+        B = 2
+        sinks = 10 * mx.random.normal(shape=(Nq,))
+        q = mx.random.normal(shape=(B, Nq, 1, D))
+        for L in (256, 8192):
+            k = mx.random.normal(shape=(B, Nkv, L, D))
+            v = mx.random.normal(shape=(B, Nkv, L, D))
+            for s_in in (None, sinks):
+                with self.subTest(L=L, sinks=s_in is not None):
+                    ref = mlx_ref_attn(q, k, v, scale, sinks=s_in)
+                    out = mx.fast.scaled_dot_product_attention(
+                        q,
+                        k,
+                        v,
+                        scale=scale,
+                        sinks=s_in,
+                        force_fused=True,
+                    )
+                    self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
+
     def test_sdpa_fully_masked(self):
         Lkv = 8
         mask = mx.array(False)
@@ -888,15 +969,21 @@ def make_qkv(qL, kL, D, qH=8, kH=8):
                 self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3))
 
         # Vector attention kernel.
-        for D in (192, 256):
+        for D in (192, 256, 512):
             with self.subTest(head_dim=D):
                 q, k, v = make_qkv(4, 16385, D, 4, 2)
                 scale = D**-0.5
-                ref = mlx_ref_attn(q, k, v, scale=scale)
-                out = mx.fast.scaled_dot_product_attention(
-                    q, k, v, scale=scale, force_fused=True
+                threshold = (
+                    patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"})
+                    if D == 512
+                    else nullcontext()
                 )
-                self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3))
+                with threshold:
+                    ref = mlx_ref_attn(q, k, v, scale=scale)
+                    out = mx.fast.scaled_dot_product_attention(
+                        q, k, v, scale=scale, force_fused=True
+                    )
+                    self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3))
 
         # No full attention fused kernels.
         with self.assertRaisesRegex(ValueError, "supports head dims"):
@@ -928,6 +1015,13 @@ def make_qkv(qL, kL, D, qH=8, kH=8):
             mx.fast.scaled_dot_product_attention(
                 q, k, v, scale=64**-0.5, force_fused=True
             )
+        with self.assertRaisesRegex(ValueError, r"requires at least \d+ keys"):
+            with patch.dict(os.environ):
+                os.environ.pop("MLX_SDPA_D512_MIN_KL", None)
+                q, k, v = make_qkv(1, 512, 512, qH=32, kH=4)
+                mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=512**-0.5, force_fused=True
+                )
 
         # No CPU fused kernel.
         with mx.stream(mx.cpu):