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..66690e75f1 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..8fbdcda2f7 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,21 +11,26 @@ "_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) \ - 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) \ - instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ - 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_helper(iname, itype, bq, bk, bd, wm, wn) \ + instantiate_attn(iname, itype, bq, bk, bd, wm, wn, iname, itype) \ + instantiate_attn(iname, itype, bq, bk, bd, wm, wn, bool_, bool) -#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_shapes_helper(iname, itype) \ + instantiate_attn_helper(iname, itype, 32, 16, 192, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 16, 128, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 96, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 80, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 72, 4, 1) \ + instantiate_attn_helper(iname, itype, 32, 32, 64, 4, 1) -instantiate_attn_mask_helper(float16, half); -instantiate_attn_mask_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float16, half); +instantiate_attn_shapes_helper(bfloat16, bfloat16_t); +instantiate_attn_shapes_helper(float32, float); + +instantiate_attn_helper(float16, half, 32, 16, 256, 4, 2); +instantiate_attn_helper(float16, half, 32, 32, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 16, 256, 4, 2); +instantiate_attn_helper(bfloat16, bfloat16_t, 32, 32, 256, 4, 2); +instantiate_attn_helper(float32, float, 32, 16, 256, 4, 2); -instantiate_attn_mask_helper(float32, float); // 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..c8307afd91 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" @@ -303,12 +303,15 @@ void sdpa_full_self_attention_metal( using namespace mlx::steel; - int wm = 4; - int wn = 1; - + char devc = d.get_architecture().back(); int bd = q.shape(-1); int bq = 32; - int bk = bd < 128 ? 32 : 16; + int bk = (bd == 256) && (q.dtype() != float32) && (devc == 'd') + ? 32 + : (bd < 128 ? 32 : 16); + + int wm = 4; + int wn = (bd == 256) ? 2 : 1; 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..18c64ed77f 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -933,23 +933,31 @@ def test_grad(slow, fast, args): 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. - for D, qL, mask in product((192, 256), (9, 16), (None, "causal")): - with self.subTest(head_dim=D, qL=qL, mask=mask): - q, k, v = make_qkv(qL, 512, D, 8, 4) + for D, qL, kL, mask, dtype in product( + (192, 256), + (9, 16), + (31, 512), + (None, "causal"), + (mx.bfloat16, mx.float32), + ): + with self.subTest(head_dim=D, qL=qL, kL=kL, mask=mask): + q, k, v = make_qkv(qL, kL, D, 8, 4) scale = D**-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 ) - self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + tol = 1e-2 if dtype == mx.bfloat16 else 1e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) # Vector attention kernel. for D in (192, 256, 512):