Use NAX attention for long unmasked D72/D80 inputs

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -231,6 +231,76 @@ void sdpa_full_self_attention_metal(
         /* const std::optional<array>& sinks = */ sinks);
   }
 
+  // Pad head dims 72 and 80 to 96 to reach the NAX kernel. The added lanes
+  // are zero and the caller's scale is retained. Enable by default only for
+  // long, unmasked half-precision attention, where the attention work can
+  // amortize the padding and output copy. The override is read per call.
+  bool pad_default = qL >= 512 && kL >= 512 && !do_causal_ && !mask && !sinks;
+  if ((D == 72 || D == 80) && metal::is_nax_available() &&
+      (q.dtype() == float16 || q.dtype() == bfloat16) &&
+      env::get_var("MLX_SDPA_PAD_HEAD_DIM", pad_default ? 1 : 0) == 1) {
+    constexpr int pad_to = 96;
+    auto& enc = metal::get_command_encoder(s);
+    array zero = array(0, q.dtype());
+
+    auto pad_head_dim = [&](const array& x) {
+      Shape padded_shape = x.shape();
+      padded_shape.back() = pad_to;
+      array xp(std::move(padded_shape), x.dtype(), nullptr, {});
+      fill_gpu(zero, xp, s);
+      // Explicit output strides: a [.., D] view over xp would be
+      // non-contiguous with span > size, so it cannot carry xp's flags.
+      copy_gpu_inplace(
+          /* const array& in = */ x,
+          /* array& out = */ xp,
+          /* const Shape& data_shape = */ x.shape(),
+          /* const Strides& i_strides = */ x.strides(),
+          /* const Strides& o_strides = */ xp.strides(),
+          /* int64_t i_offset = */ 0,
+          /* int64_t o_offset = */ 0,
+          /* CopyType ctype = */ CopyType::GeneralGeneral,
+          /* const Stream& s = */ s);
+      enc.add_temporary(xp);
+      return xp;
+    };
+
+    array qp = pad_head_dim(q);
+    array kp = pad_head_dim(k);
+    array vp = pad_head_dim(v);
+
+    Shape padded_out = o.shape();
+    padded_out.back() = pad_to;
+    array op(std::move(padded_out), o.dtype(), nullptr, {});
+    op.set_data(allocator::malloc(op.nbytes()));
+
+    sdpa_full_self_attention_nax(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& q = */ qp,
+        /* const array& k = */ kp,
+        /* const array& v = */ vp,
+        /* const float scale = */ scale,
+        /* array& o = */ op,
+        /* bool do_causal_ = */ do_causal_,
+        /* const std::optional<array>& mask = */ mask,
+        /* const std::optional<array>& sinks = */ sinks);
+
+    // Slice the padded lanes back out into o's caller-chosen strides.
+    copy_gpu_inplace(
+        /* const array& in = */ op,
+        /* array& out = */ o,
+        /* const Shape& data_shape = */ o.shape(),
+        /* const Strides& i_strides = */ op.strides(),
+        /* const Strides& o_strides = */ o.strides(),
+        /* int64_t i_offset = */ 0,
+        /* int64_t o_offset = */ 0,
+        /* CopyType ctype = */ CopyType::GeneralGeneral,
+        /* const Stream& s = */ s);
+    enc.add_temporary(op);
+    enc.add_temporary(zero);
+    return;
+  }
+
   using namespace mlx::steel;
 
   int wm = 4;

python/tests/test_fast_sdpa.py

@@ -144,6 +144,68 @@ def test_sdpa_head_dim_72(self):
                 diff = mx.abs(out - ref) - atol * mx.abs(ref)
                 self.assertLessEqual(mx.max(diff).item(), atol)
 
+    @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only")
+    def test_sdpa_pad_head_dim_opt_in(self):
+        if mx.default_device() != mx.gpu:
+            self.skipTest("requires GPU")
+        name = "MLX_SDPA_PAD_HEAD_DIM"
+        previous = os.environ.get(name)
+        try:
+            os.environ[name] = "1"
+            self.test_sdpa_head_dim_72()
+            self.test_sdpa_head_dim_80()
+            self.test_sdpa_head_dim_72_80_sinks()
+        finally:
+            if previous is None:
+                os.environ.pop(name, None)
+            else:
+                os.environ[name] = previous
+
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only")
+    def test_sdpa_head_dim_80(self):
+        B, D, qH, kH = (1, 80, 8, 2)
+        for qL, kL, dtype, mask_str in product(
+            (64, 65),
+            (128, 127),
+            (mx.float16, mx.bfloat16, mx.float32),
+            (None, "additive", "bool", "causal"),
+        ):
+            with self.subTest(qL=qL, kL=kL, dtype=dtype, mask=mask_str):
+                q, k, v, scale, mask = prepare_inputs(
+                    B, qL, kL, D, qH, kH, mask_str, False, dtype
+                )
+                ref = mlx_ref_attn(q, k, v, scale, mask)
+                out = mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=scale, mask=mask
+                )
+
+                if dtype == mx.float32:
+                    atol = 1e-5
+                elif dtype == mx.bfloat16:
+                    atol = 5e-3
+                else:
+                    atol = 3e-4
+                diff = mx.abs(out - ref) - atol * mx.abs(ref)
+                self.assertLessEqual(mx.max(diff).item(), atol)
+
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only")
+    def test_sdpa_head_dim_72_80_sinks(self):
+        B, qH, kH, qL, kL = (1, 8, 2, 64, 128)
+        for D, dtype in product((72, 80), (mx.float16, mx.bfloat16)):
+            with self.subTest(D=D, dtype=dtype):
+                q, k, v, scale, _ = prepare_inputs(
+                    B, qL, kL, D, qH, kH, None, False, dtype
+                )
+                sinks = 10 * mx.random.normal(shape=(qH,), dtype=dtype)
+                ref = mlx_ref_attn(q, k, v, scale, sinks=sinks)
+                out = mx.fast.scaled_dot_product_attention(
+                    q, k, v, scale=scale, sinks=sinks
+                )
+
+                atol = 5e-3 if dtype == mx.bfloat16 else 3e-4
+                diff = mx.abs(out - ref) - atol * mx.abs(ref)
+                self.assertLessEqual(mx.max(diff).item(), atol)
+
     @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only")
     def test_sdpa_head_dim_96(self):
         B, D, qH, kH = (1, 96, 8, 2)