diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h index 29fa7ba396..614c3bb3ea 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h @@ -1,4 +1,4 @@ -// Copyright © 2024-25 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include "mlx/backend/metal/kernels/steel/attn/attn.h" @@ -119,8 +119,10 @@ template < constexpr short tgp_mem_0 = (BK + padK) * (BD); constexpr short tgp_mem_1 = BK * (BD + padV); constexpr short tgp_mem_s = tgp_mem_0 > tgp_mem_1 ? tgp_mem_0 : tgp_mem_1; + constexpr bool q_in_registers = WN == 2 && (BK == 32 || sizeof(T) == 4); + constexpr short q_tgp_mem_s = q_in_registers ? 1 : BQ * (BD + padQ); - threadgroup T Q_smem[BQ * (BD + padQ)]; + threadgroup T Q_smem[q_tgp_mem_s]; threadgroup T KV_smem[tgp_mem_s]; threadgroup T* Qs = Q_smem; @@ -168,8 +170,13 @@ template < // Prepare MMA tiles constexpr short kFragSize = 8; // MMAFrag size using MMAFrag_acc_t = BaseMMAFrag; + using MMAInType = metal::conditional_t; + using MMAFrag_in_t = BaseMMAFrag; - constexpr int kNWarps = WM * WN; + static_assert(WN == 1 || WN == 2, "WN must be 1 or 2"); + static_assert(BD % WN == 0, "The head dim must split evenly across WN"); + + constexpr int kNWarps = WM; static_assert( BQ >= (kNWarps * kFragSize) && BQ % (kNWarps * kFragSize) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); @@ -180,14 +187,17 @@ template < constexpr int TK = BK / kFragSize; // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kFragSize; + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; static_assert(TQ == 1, "Check TQ"); - MMATile Qtile; - MMATile Ktile; + MMATile Qtile; + MMATile Ktile; MMATile Stile; - MMATile Vtile; - MMATile Otile; + MMATile Vtile; + MMATile Otile; + MMATile Qtiles[q_in_registers ? TDh : 1]; Otile.clear(); @@ -195,22 +205,45 @@ template < const short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = kFragSize * TQ * simd_group_id; + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + const short tm = kFragSize * TQ * row_group; - const short Qs_offset = (tm + sm) * LDQ_tgp + sn; - const short Ks_offset = sm * LDK_tgp + sn; - const short Vs_offset = sm * LDV_tgp + sn; + const short Qs_offset = (tm + sm) * LDQ_tgp + d_half * BDh + sn; + const short Ks_offset = d_half * BDh * LDK_tgp + sm * LDK_tgp + sn; + const short Vs_offset = sm * LDV_tgp + d_half * BDh + sn; constexpr short Qs_tile_stride = kFragSize; constexpr short Ks_tile_stride = kFragSize * LDK_tgp; threadgroup_barrier(mem_flags::mem_threadgroup); - // Load Q blocks - if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - loader_q.load_safe(short2(BD, params->qL_rem)); + if constexpr (!q_in_registers) { + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + loader_q.load_safe(short2(BD, params->qL_rem)); + } else { + loader_q.load_unsafe(); + } } else { - loader_q.load_unsafe(); + STEEL_PRAGMA_UNROLL + for (short dd = 0; dd < TDh; dd++) { + const device T* q_src = Q + (tm + sm) * params->Q_strides[2] + + d_half * BDh + dd * kFragSize + sn; + if (!align_Q && int(tid.x) == (params->NQ_aligned)) { + MMAFrag_in_t::load_safe( + Qtiles[dd].frag_at(0, 0), + q_src, + params->Q_strides[2], + Int<1>{}, + params->qL_rem - (tm + sm), + BDh - (dd * kFragSize + sn), + Int<0>{}, + Int<0>{}); + } else { + MMAFrag_in_t::load( + Qtiles[dd].frag_at(0, 0), q_src, params->Q_strides[2], Int<1>{}); + } + } } // Init row reduction variables @@ -262,17 +295,51 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); STEEL_PRAGMA_UNROLL - for (short dd = 0; dd < TD; dd++) { + for (short dd = 0; dd < TDh; dd++) { simdgroup_barrier(mem_flags::mem_none); - Qtile.template load( - &Qs[Qs_offset + dd * Qs_tile_stride]); + if constexpr (!q_in_registers) { + Qtile.template load( + &Qs[Qs_offset + dd * Qs_tile_stride]); + } Ktile.template load( &Ks[Ks_offset + dd * Ks_tile_stride]); simdgroup_barrier(mem_flags::mem_none); - tile_matmad(Stile, Qtile, Ktile, Stile); + if constexpr (!q_in_registers) { + tile_matmad(Stile, Qtile, Ktile, Stile); + } else { + tile_matmad(Stile, Qtiles[dd], Ktile, Stile); + } + } + + if constexpr (WN == 2) { + constexpr int s_xchg_stride = decltype(Stile)::kElemsPerTile * 32; + constexpr int s_xchg_size = WM * WN * s_xchg_stride; + static_assert( + s_xchg_size * sizeof(AccumType) <= tgp_mem_s * sizeof(T), + "The score exchange must fit in KV_smem"); + + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup AccumType* s_xchg = + reinterpret_cast(KV_smem); + threadgroup AccumType* slot = + s_xchg + (row_group * WN + d_half) * s_xchg_stride; + const short base = short(simd_lane_id) * decltype(Stile)::kElemsPerTile; + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + slot[base + ii] = Stile.elems()[ii]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const threadgroup AccumType* peer = + s_xchg + (row_group * WN + 1 - d_half) * s_xchg_stride; + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + Stile.elems()[ii] += peer[base + ii]; + } } // Apply scale in float32 @@ -425,7 +492,7 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TDh; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { if constexpr (BD >= 128) { @@ -461,10 +528,10 @@ template < threadgroup_barrier(mem_flags::mem_none); // Store results - O += (tm + sm) * params->O_strides[2] + sn; + O += (tm + sm) * params->O_strides[2] + d_half * BDh + sn; if (!align_Q && int(tid.x) == (params->NQ_aligned)) { - auto dst_tile_dims = short2(BD - sn, params->qL_rem - (tm + sm)); + auto dst_tile_dims = short2(BDh - sn, params->qL_rem - (tm + sm)); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) return; diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal index fbd84004f0..a2d772991f 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal +++ b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.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" @@ -11,8 +11,8 @@ "_wm" #wm "_wn" #wn "_mask" #mname, \ attention, dtype, bq, bk, bd, wm, wn, mtype, float) -#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ - instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \ +#define instantiate_attn_shapes_helper(iname, itype, mname, mtype, bk256) \ + instantiate_attn(iname, itype, 32, bk256, 256, 4, 2, mname, mtype) \ instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ @@ -20,12 +20,16 @@ instantiate_attn(iname, itype, 32, 32, 72, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype) -#define instantiate_attn_mask_helper(iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, bool_, bool) +#define instantiate_attn_mask_helper(iname, itype, bk256) \ + instantiate_attn_shapes_helper(iname, itype, iname, itype, bk256) \ + instantiate_attn_shapes_helper(iname, itype, bool_, bool, bk256) -instantiate_attn_mask_helper(float16, half); -instantiate_attn_mask_helper(bfloat16, bfloat16_t); +instantiate_attn_mask_helper(float16, half, 32); +instantiate_attn_mask_helper(bfloat16, bfloat16_t, 32); +instantiate_attn(float16, half, 32, 16, 256, 4, 2, float16, half); +instantiate_attn(float16, half, 32, 16, 256, 4, 2, bool_, bool); +instantiate_attn(bfloat16, bfloat16_t, 32, 16, 256, 4, 2, bfloat16, bfloat16_t); +instantiate_attn(bfloat16, bfloat16_t, 32, 16, 256, 4, 2, bool_, bool); -instantiate_attn_mask_helper(float32, float); +instantiate_attn_mask_helper(float32, float, 16); // clang-format on diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 6ece8c43fb..b8f0838a4c 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -1,4 +1,4 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-26 Apple Inc. #include #include "mlx/backend/common/compiled.h" @@ -304,11 +304,14 @@ void sdpa_full_self_attention_metal( using namespace mlx::steel; int wm = 4; - int wn = 1; int bd = q.shape(-1); + char devc = d.get_architecture().back(); + int wn = bd == 256 ? 2 : 1; int bq = 32; - int bk = bd < 128 ? 32 : 16; + int bk = bd == 256 && q.dtype() != float32 && devc == 'd' + ? 32 + : (bd < 128 ? 32 : 16); const bool align_Q = (qL % bq) == 0; const bool align_K = (kL % bk) == 0; @@ -862,6 +865,12 @@ bool ScaledDotProductAttention::use_fallback( return false; } + if (!metal::is_nax_available() && q.dtype() != float32 && do_causal && + query_head_dim == 256 && query_sequence_length >= 2048 && + query_sequence_length == k.shape(2)) { + return false; + } + // Unfused path is faster for following shapes. if (query_sequence_length > 8) { return query_head_dim == 192 || query_head_dim == 256; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index fefe7325b3..e8daa1fb86 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -929,15 +929,35 @@ def test_grad(slow, fast, args): ).sum() test_grad(loss_slow, loss_fast, [q, k, v]) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_d256_default_metal(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + mx.random.seed(0) + + for dtype in (mx.float16, mx.bfloat16): + with self.subTest(dtype=dtype): + q = (0.5 * mx.random.normal((1, 2, 2048, 256))).astype(dtype) + k = (0.5 * mx.random.normal((1, 1, 2048, 256))).astype(dtype) + v = (0.5 * mx.random.normal((1, 1, 2048, 256))).astype(dtype) + scale = 256**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale, mask="causal") + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask="causal" + ) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-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_force_fused_metal(self): if mx.default_device() != mx.gpu: self.skipTest("requires GPU") + mx.random.seed(0) - def make_qkv(qL, kL, D, qH=8, kH=8): - q = mx.random.normal((1, qH, qL, D), mx.float16) - k = mx.random.normal((1, kH, kL, D), mx.float16) - v = mx.random.normal((1, kH, kL, D), mx.float16) + def make_qkv(qL, kL, D, qH=8, kH=8, dtype=mx.float16): + q = mx.random.normal((1, qH, qL, D), dtype) + k = mx.random.normal((1, kH, kL, D), dtype) + v = mx.random.normal((1, kH, kL, D), dtype) return q, k, v # Full attention kernel. @@ -951,6 +971,33 @@ def make_qkv(qL, kL, D, qH=8, kH=8): ) self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + for dtype, mask_kind in product( + (mx.bfloat16, mx.float32), (None, "causal", "bool", "additive") + ): + with self.subTest(head_dim=256, dtype=dtype, mask=mask_kind): + q, k, v = make_qkv(9, 31, 256, 8, 4, dtype) + if mask_kind == "bool": + mask = mx.random.uniform(shape=(1, 8, 9, 31)) > 0.2 + elif mask_kind == "additive": + mask = mx.random.normal((1, 8, 9, 31), dtype) + else: + mask = mask_kind + scale = 256**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask, force_fused=True + ) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + + q, k, _ = make_qkv(9, 31, 256, 8, 4, mx.bfloat16) + v_half = mx.random.normal((1, 4, 31, 128), mx.bfloat16) + v = mx.concatenate((v_half, v_half), axis=-1) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=256**-0.5, force_fused=True + ) + self.assertTrue(mx.array_equal(out[..., :128], out[..., 128:])) + # Vector attention kernel. for D in (192, 256, 512): with self.subTest(head_dim=D):