diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index 4a5a9716fd..6d72d56b58 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/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; 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; 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 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 new_max; metal::vec 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 diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal index 66d55539ab..08d703df3f 100644 --- a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal +++ b/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 diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 6ece8c43fb..9b2abffa91 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" @@ -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 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() && diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index fefe7325b3..556923dba5 100644 --- a/python/tests/test_fast_sdpa.py +++ b/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"