From 69e88a18a4ed67a4a7a9a2418a9a9cad4fef9be4 Mon Sep 17 00:00:00 2001
From: Yanzhao Wang <yzwang1996@gmail.com>
Date: Tue, 8 Sep 2026 00:38:14 -0700
Subject: [PATCH] Use NAX attention for short causal D256 prefill

---
 .../metal/scaled_dot_product_attention.cpp    | 17 +++++--
 python/tests/test_fast_sdpa.py                | 49 +++++++++++++++++++
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp
index 6d9ade1ccc..1c8d6b873c 100644
--- a/mlx/backend/metal/scaled_dot_product_attention.cpp
+++ b/mlx/backend/metal/scaled_dot_product_attention.cpp
@@ -823,13 +823,20 @@ bool ScaledDotProductAttention::use_fallback(
   const int query_head_dim = q.shape(-1);
   const int value_head_dim = v.shape(-1);
 
-  // Use headdim-split kernel when NAX is enabled and there are enough query
-  // blocks to fill the machine.
+  // Use the head-dim-split kernel from 1024 query rows. For fp16/bf16 also
+  // admit short causal prefill chunks while the key length stays at most
+  // 1536: a 512-query causal chunk is faster on that kernel there, while
+  // float32 and explicit array masks are not, so they keep the routing below.
   if (metal::is_nax_available() &&
-      (env::enable_tf32() || q.dtype() != float32) &&
-      query_sequence_length >= 1024 && query_head_dim == 256 &&
+      (env::enable_tf32() || q.dtype() != float32) && query_head_dim == 256 &&
       (do_causal || has_arr_mask)) {
-    return false;
+    if (query_sequence_length >= 1024) {
+      return false;
+    }
+    if (do_causal && (q.dtype() == float16 || q.dtype() == bfloat16) &&
+        query_sequence_length >= 512 && k.shape(2) <= 1536) {
+      return false;
+    }
   }
 
   // Unfused path is faster for following shapes.
diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py
index 6b185ecf1f..ad3139e821 100644
--- a/python/tests/test_fast_sdpa.py
+++ b/python/tests/test_fast_sdpa.py
@@ -300,6 +300,55 @@ def test_sdpa_full_head_dim_256(self):
                         tol = 5e-3
                     self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol))
 
+    # The shapes cover the causal D=256 window (512 <= qL < 1024,
+    # kL <= 1536) and its kL edge.
+    @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
+    def test_sdpa_head_dim_256_causal_prefill_window(self):
+        if mx.default_device() != mx.gpu:
+            self.skipTest("requires GPU")
+        D, Nq, Nkv = 256, 16, 2
+        scale = D**-0.5
+        for dtype, (qL, kL) in product(
+            (mx.float16, mx.bfloat16),
+            ((512, 512), (512, 1536), (512, 1537), (768, 1024), (1023, 1536)),
+        ):
+            with self.subTest(dtype=dtype, qL=qL, kL=kL):
+                mx.random.seed(0)
+                q = (0.5 * mx.random.normal((1, Nq, qL, D))).astype(dtype)
+                k = (0.5 * mx.random.normal((1, Nkv, kL, D))).astype(dtype)
+                v = (0.5 * mx.random.normal((1, Nkv, kL, D))).astype(dtype)
+                ref = mlx_primitives_sdpa(
+                    q,
+                    mx.repeat(k, Nq // Nkv, axis=1),
+                    mx.repeat(v, Nq // Nkv, axis=1),
+                    scale,
+                    mask="causal",
+                )
+                out = mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=scale, mask="causal"
+                )
+                self.assertTrue(mx.allclose(ref, out, atol=5e-3, rtol=5e-3))
+
+        # Exercise batch/head strides, an odd tail and sinks on the new path.
+        for dtype in (mx.float16, mx.bfloat16):
+            with self.subTest(dtype=dtype, sinks=True):
+                q = mx.random.normal((2, 6, 513, D)).astype(dtype)
+                k = mx.random.normal((2, 2, 1057, D)).astype(dtype)[:, :, 1:1026]
+                v = mx.random.normal((2, 2, 1057, D)).astype(dtype)[:, :, 1:1026]
+                sinks = mx.linspace(5, 10, 6).astype(dtype)
+                ref = mlx_ref_attn(
+                    q.astype(mx.float32),
+                    k.astype(mx.float32),
+                    v.astype(mx.float32),
+                    scale,
+                    mask="causal",
+                    sinks=sinks.astype(mx.float32),
+                )
+                out = mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=scale, mask="causal", sinks=sinks
+                )
+                self.assertTrue(mx.allclose(ref, out, atol=5e-3, rtol=5e-3))
+
     def test_sdpa_vector_kv_transposed_head_seq(self):
         D = 64
         Nq = 4