diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 286535e124..6d9ade1ccc 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -231,6 +231,76 @@ void sdpa_full_self_attention_metal( /* const std::optional& 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& mask = */ mask, + /* const std::optional& 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; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 1d2987e6a1..6b185ecf1f 100644 --- a/python/tests/test_fast_sdpa.py +++ b/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)