Improve metal memory usage for SDPA D512

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

@@ -1,4 +1,4 @@
-// Copyright © 2024-25 Apple Inc.
+// Copyright © 2024-26 Apple Inc.
 
 #include "mlx/backend/metal/kernels/steel/attn/nax.h"
 #include "mlx/backend/metal/kernels/steel/attn/params.h"
@@ -489,15 +489,10 @@ template <
 // Head-dim split attention kernel
 ///////////////////////////////////////////////////////////////////////////////
 
-// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup
-// accumulator working set of attention_nax (TD output fragments plus the S
-// fragments) is what gates tensor-unit throughput, so this kernel splits the
-// head dim across the WN = 2 simdgroups of the second warp dimension: each
-// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V,
-// halving its accumulator set. The pair exchanges its partial Q@K.T sums
-// through threadgroup memory, then both simdgroups run softmax redundantly on
-// the full S tile (the row statistics are cheap) and each accumulates P@V for
-// its own half of Dv.
+// Split wide heads across WN simdgroups to reduce the accumulator working set.
+// Each group owns one D slice. The groups exchange their partial Q@K.T scores
+// through threadgroup memory and reduce them in the same order. Each group then
+// runs softmax and accumulates its P@V slice.
 
 // clang-format off
 template <
@@ -556,9 +551,9 @@ template <
   constexpr short kU = 16;
 
   // The WM simdgroups along the first warp dimension split the Q sequence;
-  // the WN simdgroups along the second split the head dim. The exchange
-  // below reduces exactly one peer, so WN is fixed at 2.
-  static_assert(WN == 2, "The head-dim split kernel needs WN == 2");
+  // the WN simdgroups along the second split the head dim.
+  static_assert(
+      WN == 2 || WN == 4, "The head-dim split kernel needs WN == 2 or WN == 4");
   constexpr int kNWarps = WM;
   static_assert(
       BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0,
@@ -579,20 +574,20 @@ template <
   constexpr int BDh = BD / WN;
 
   static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs");
-  static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair");
+  static_assert(TK % 2 == 0, "S fragments are computed pair by pair");
 
   const short row_group = simd_group_id / WN;
-  const short d_half = simd_group_id % WN;
+  const short d_group = simd_group_id % WN;
 
   using otile_t = NAXTile<AccumType, TQ, TDh>;
   otile_t Otile;
   Otile.clear();
 
   const short tm = kU * TQ * row_group;
-  Q += tm * int(params->Q_strides[2]) + d_half * BDh;
-  K += d_half * BDh;
-  V += d_half * BDh;
-  O += tm * int(params->O_strides[2]) + d_half * BDh;
+  Q += tm * int(params->Q_strides[2]) + d_group * BDh;
+  K += d_group * BDh;
+  V += d_group * BDh;
+  O += tm * int(params->O_strides[2]) + d_group * BDh;
 
   constexpr short kRowsPT = otile_t::kRowsPerThread;
 
@@ -632,12 +627,11 @@ template <
   using stile_t = NAXTile<AccumType, TQ, TK>;
   constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag;
 
-  // One slot per (row group, half): a fragment pair in per-lane-linear
-  // layout. Both halves share the fragment-to-lane mapping, so the
-  // exchange needs no coordinate math.
-  threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32];
+  // One slot per (row group, D group) in per-lane-linear layout. All groups
+  // share the fragment-to-lane mapping.
+  threadgroup AccumType s_xchg[WM][WN][TK * kEPF * 32];
 
-  // Keep the simdgroup's Q half resident in registers for the whole KV
+  // Keep the simdgroup's Q slice resident in registers for the whole KV
   // loop: TDh fragments of T are cheap next to the accumulators.
   NAXTile<T, 1, 1> Qtiles[TDh];
   STEEL_PRAGMA_UNROLL
@@ -662,7 +656,7 @@ template <
     stile_t Stile;
     Stile.clear();
 
-    // S = Q @ K.T, this half of D only, exchanged pair by pair.
+    // S = Q @ K.T for this D slice.
     STEEL_PRAGMA_UNROLL
     for (short ik = 0; ik < TK; ik += 2) {
       STEEL_PRAGMA_UNROLL
@@ -686,26 +680,37 @@ template <
             Ktile.frag_at(1, 0),
             metal::true_type{});
       }
+    }
 
-      // Exchange the partial pair and reduce.
-      threadgroup AccumType* slot = s_xchg[row_group][d_half];
-      thread auto& s0 = Stile.frag_at(0, ik);
-      thread auto& s1 = Stile.frag_at(0, ik + 1);
-      const short base = short(simd_lane_id) * (2 * kEPF);
+    // Exchange all partial scores and reduce them in a fixed order.
+    threadgroup AccumType* slot = s_xchg[row_group][d_group];
+    const short base = short(simd_lane_id) * (TK * kEPF);
+    STEEL_PRAGMA_UNROLL
+    for (short ik = 0; ik < TK; ik++) {
+      thread auto& s = Stile.frag_at(0, ik);
       STEEL_PRAGMA_UNROLL
       for (short i = 0; i < kEPF; i++) {
-        slot[base + i] = s0[i];
-        slot[base + kEPF + i] = s1[i];
+        slot[base + ik * kEPF + i] = s[i];
       }
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half];
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    STEEL_PRAGMA_UNROLL
+    for (short ik = 0; ik < TK; ik++) {
+      thread auto& s = Stile.frag_at(0, ik);
       STEEL_PRAGMA_UNROLL
       for (short i = 0; i < kEPF; i++) {
-        s0[i] += peer[base + i];
-        s1[i] += peer[base + kEPF + i];
+        s[i] = s_xchg[row_group][0][base + ik * kEPF + i];
+      }
+      STEEL_PRAGMA_UNROLL
+      for (short peer_group = 1; peer_group < WN; peer_group++) {
+        const threadgroup AccumType* peer = s_xchg[row_group][peer_group];
+        STEEL_PRAGMA_UNROLL
+        for (short i = 0; i < kEPF; i++) {
+          s[i] += peer[base + ik * kEPF + i];
+        }
       }
-      threadgroup_barrier(mem_flags::mem_threadgroup);
     }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
 
     // Scale S
     STEEL_PRAGMA_UNROLL
@@ -826,7 +831,7 @@ template <
       }
     }
 
-    // Do softmax (redundantly per half; the row statistics are cheap)
+    // Do softmax in each D group; the row statistics are cheap.
     metal::vec<AccumType, kRowsPT> new_max;
     metal::vec<AccumType, kRowsPT> factor;
     STEEL_PRAGMA_UNROLL
@@ -854,7 +859,7 @@ template <
 
     simdgroup_barrier(mem_flags::mem_none);
 
-    // O = P @ V, this half of Dv only.
+    // O = P @ V for this D slice.
     STEEL_PRAGMA_UNROLL
     for (short id = 0; id < TDh; id += 2) {
       STEEL_PRAGMA_UNROLL

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

@@ -1,4 +1,4 @@
-// Copyright © 2024-25 Apple Inc.
+// Copyright © 2024-26 Apple Inc.
 
 // clang-format off
 #include "mlx/backend/metal/kernels/utils.h"
@@ -18,6 +18,7 @@
   attention_nax_dsplit, dtype, bq, bk, bd, wm, wn, mtype, float)
 
 #define instantiate_attn_shapes_helper(iname, itype, mname, mtype)         \
+    instantiate_attn_dsplit(iname, itype, 32, 32, 512, 2, 4, mname, mtype) \
     instantiate_attn_dsplit(iname, itype, 64, 32, 256, 4, 2, mname, mtype) \
     instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype)        \
     instantiate_attn(iname, itype, 64, 32,  96, 4, 1, mname, mtype)        \
@@ -31,6 +32,5 @@
 
 instantiate_attn_mask_helper(float16, half);
 instantiate_attn_mask_helper(bfloat16, bfloat);
-
 instantiate_attn_mask_helper(float32, float);
 // clang-format on

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2024 Apple Inc.
+// Copyright © 2024-26 Apple Inc.
 #include <sstream>
 
 #include "mlx/backend/common/compiled.h"
@@ -29,13 +29,12 @@ void sdpa_full_self_attention_nax(
   using namespace mlx::steel;
 
   int bd = q.shape(-1);
-  int bq = 64;
+  int bq = bd == 512 ? 32 : 64;
   int bk = 32;
 
-  bool split_d = bd == 256;
-  int wm = 4;
-  int wn = split_d ? 2 : 1;
-
+  bool split_d = bd == 256 || bd == 512;
+  int wm = bd == 512 ? 2 : 4;
+  int wn = split_d ? bd / 128 : 1;
   int B = q.shape(0);
   int H = q.shape(1);
   int D = q.shape(3);
@@ -216,7 +215,7 @@ void sdpa_full_self_attention_metal(
   int kL = k.shape(2);
 
   if (metal::is_nax_available() &&
-      (D == 64 || D == 96 || D == 128 || D == 256) &&
+      (D == 64 || D == 96 || D == 128 || D == 256 || D == 512) &&
       (env::enable_tf32() || q.dtype() != float32)) {
     return sdpa_full_self_attention_nax(
         /* const Stream& s = */ s,
@@ -733,15 +732,19 @@ std::tuple<bool, std::string> has_fused_kernel(
 
   std::ostringstream msg;
   if (query_sequence_length > 8) {
+    const bool supports_d512 = metal::is_nax_available() &&
+        (env::enable_tf32() || q.dtype() != float32);
     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);
+         query_head_dim == 192 || query_head_dim == 256 ||
+         (query_head_dim == 512 && supports_d512));
     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 << ".";
+          << "{64, 72, 80, 96, 128, 192, 256} with matching query/value dims, "
+          << "plus head dim 512 on NAX GPUs (float32 also requires TF32); got "
+          << "query head dim " << query_head_dim << " and value head dim "
+          << value_head_dim << ".";
       return {false, msg.str()};
     }
     if (has_mask && !has_arr_mask &&
@@ -853,6 +856,17 @@ bool ScaledDotProductAttention::use_fallback(
   const int query_head_dim = q.shape(-1);
   const int value_head_dim = v.shape(-1);
 
+  if (query_head_dim == 512 && query_sequence_length > 8) {
+    constexpr int64_t min_query_blocks = 1024;
+    const int64_t query_blocks = int64_t(q.shape(0)) * q.shape(1) *
+        int64_t(ceildiv(query_sequence_length, 32));
+    // The D512 kernel needs this many query blocks to match fallback speed.
+    const bool eligible = metal::is_nax_available() && q.dtype() != float32 &&
+        query_sequence_length >= 1024 && do_causal && !has_arr_mask &&
+        query_blocks >= min_query_blocks;
+    return !eligible;
+  }
+
   // Use headdim-split kernel when NAX is enabled and there are enough query
   // blocks to fill the machine.
   if (metal::is_nax_available() &&

python/tests/test_fast_sdpa.py

@@ -292,6 +292,120 @@ 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_head_dim_512_nax(self):
+        if mx.default_device() != mx.gpu:
+            self.skipTest("requires GPU")
+
+        D = 512
+        scale = D**-0.5
+        mx.random.seed(0)
+
+        q_probe = mx.random.normal(shape=(1, 2, 9, D)).astype(mx.bfloat16)
+        k_probe = mx.random.normal(shape=(1, 1, 31, D)).astype(mx.bfloat16)
+        v_probe = mx.random.normal(shape=(1, 1, 31, D)).astype(mx.bfloat16)
+        try:
+            mx.eval(
+                mx.fast.scaled_dot_product_attention(
+                    q_probe,
+                    k_probe,
+                    v_probe,
+                    scale=scale,
+                    force_fused=True,
+                )
+            )
+        except ValueError as e:
+            if "plus head dim 512 on NAX GPUs" not in str(e):
+                raise
+            self.skipTest("D512 full attention requires NAX")
+
+        cases = [
+            # Multiple query tiles with GQA.
+            (1, 8, 2, 257, 513, None, True),
+            # Multiple batches and a ragged final tile.
+            (2, 4, 1, 513, 777, None, True),
+            # Gemma 4 31B default route at the work threshold with sliced K/V.
+            (1, 32, 16, 1024, 1025, 1056, False),
+        ]
+        for dtype in (mx.float16, mx.bfloat16, mx.float32):
+            for B, Nq, Nkv, qL, kL, cache_len, force_fused in cases:
+                with self.subTest(
+                    dtype=dtype,
+                    B=B,
+                    Nq=Nq,
+                    Nkv=Nkv,
+                    qL=qL,
+                    kL=kL,
+                    cache_len=cache_len,
+                ):
+                    q = mx.random.normal(shape=(B, Nq, qL, D)).astype(dtype)
+                    if cache_len is None:
+                        k = mx.random.normal(shape=(B, Nkv, kL, D)).astype(dtype)
+                        v = mx.random.normal(shape=(B, Nkv, kL, D)).astype(dtype)
+                    else:
+                        k = mx.random.normal(shape=(B, Nkv, cache_len, D)).astype(
+                            dtype
+                        )[..., :kL, :]
+                        v = mx.random.normal(shape=(B, Nkv, cache_len, D)).astype(
+                            dtype
+                        )[..., :kL, :]
+
+                    ref = mlx_ref_attn(q, k, v, scale, "causal")
+                    out = mx.fast.scaled_dot_product_attention(
+                        q,
+                        k,
+                        v,
+                        scale=scale,
+                        mask="causal",
+                        force_fused=force_fused
+                        and (
+                            dtype != mx.float32
+                            or os.getenv("MLX_ENABLE_TF32", "1") != "0"
+                        ),
+                    )
+                    tol = 2e-3 if dtype != mx.bfloat16 else 1e-2
+                    self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol))
+
+        q = mx.random.normal(shape=(1, 8, 33, D)).astype(mx.bfloat16)
+        k = mx.random.normal(shape=(1, 2, 67, D)).astype(mx.bfloat16)
+        v = mx.random.normal(shape=(1, 2, 67, D)).astype(mx.bfloat16)
+        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-2, rtol=1e-2))
+
+        bool_mask = mx.random.uniform(shape=(1, 1, 33, 67)) > 0.2
+        additive_mask = mx.where(bool_mask, 0.0, -1e4).astype(mx.bfloat16)
+        for mask in (bool_mask, additive_mask):
+            ref = mlx_ref_attn(q, k, v, scale, mask)
+            out = mx.fast.scaled_dot_product_attention(
+                q, k, v, scale=scale, mask=mask, force_fused=True
+            )
+            self.assertTrue(mx.allclose(ref, out, atol=1e-2, rtol=1e-2))
+
+        q_same = mx.random.normal(shape=(1, 8, 33, D)).astype(mx.bfloat16)
+        k_same = mx.random.normal(shape=(1, 2, 67, D)).astype(mx.bfloat16)
+        v_slice = mx.random.normal(shape=(1, 2, 67, D // 4)).astype(mx.bfloat16)
+        v_same = mx.tile(v_slice, (1, 1, 1, 4))
+        out_same = mx.fast.scaled_dot_product_attention(
+            q_same, k_same, v_same, scale=scale, force_fused=True
+        )
+        for d_group in range(1, 4):
+            self.assertTrue(
+                mx.array_equal(
+                    out_same[..., : D // 4],
+                    out_same[..., d_group * D // 4 : (d_group + 1) * D // 4],
+                )
+            )
+
+        sinks = mx.random.normal(shape=(8,)).astype(mx.bfloat16)
+        ref = mlx_ref_attn(q, k, v, scale, sinks=sinks)
+        out = mx.fast.scaled_dot_product_attention(
+            q, k, v, scale=scale, sinks=sinks, force_fused=True
+        )
+        self.assertTrue(mx.allclose(ref, out, atol=1e-2, rtol=1e-2))
+
     def test_sdpa_vector_kv_transposed_head_seq(self):
         D = 64
         Nq = 4
@@ -965,9 +1079,9 @@ def make_qkv(qL, kL, D, qH=8, kH=8):
 
         # No full attention fused kernels.
         with self.assertRaisesRegex(ValueError, "supports head dims"):
-            q, k, v = make_qkv(16, 512, 512)
+            q, k, v = make_qkv(16, 512, 1024)
             mx.fast.scaled_dot_product_attention(
-                q, k, v, scale=512**-0.5, force_fused=True
+                q, k, v, scale=1024**-0.5, force_fused=True
             )
         with self.assertRaisesRegex(
             ValueError, "query sequence to be no longer than the key sequence"