diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 807a6550b7..1f943c96dd 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -1360,6 +1360,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bq, int bk, int bd, + int bv, int wm, int wn, const array& m, @@ -1371,16 +1372,28 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( kernel_source, metal::utils(), metal::steel_attention_nax(), - get_template_definition( - lib_name, - split_d ? "attention_nax_dsplit" : "attention_nax", - get_type_string(q.dtype()), - bq, - bk, - bd, - wm, - wn, - get_type_string(m.dtype()))); + split_d ? get_template_definition( + lib_name, + "attention_nax_dsplit", + get_type_string(q.dtype()), + bq, + bk, + bd, + wm, + wn, + get_type_string(m.dtype())) + : get_template_definition( + lib_name, + "attention_nax", + get_type_string(q.dtype()), + bq, + bk, + bd, + wm, + wn, + get_type_string(m.dtype()), + "float", + bv)); return kernel_source; }); return d.get_kernel(kernel_name, lib, hash_name, func_consts); diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index 42888f0a78..9b74726507 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -426,6 +426,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bq, int bk, int bd, + int bv, int wm, int wn, const array& m, diff --git a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index 68e2a6c06a..04206c9404 100644 --- a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal @@ -40,6 +40,7 @@ using namespace metal; #define instantiate_sdpa_vector_heads(type) \ instantiate_sdpa_vector(type, 64, 64) \ + instantiate_sdpa_vector(type, 96, 64) \ instantiate_sdpa_vector(type, 96, 96) \ instantiate_sdpa_vector(type, 128, 128) \ instantiate_sdpa_vector(type, 192, 128) \ 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..b248ffdd07 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 @@ -79,7 +79,8 @@ template < int WM, int WN, typename MaskType = float, - typename AccumType = float> + typename AccumType = float, + int BV = BD> [[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention_nax( const device T* Q [[buffer(0)]], const device T* K [[buffer(1)]], @@ -136,11 +137,13 @@ template < constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) constexpr int TD = BD / kU; + constexpr int TV = BV / kU; + static_assert(BD % kU == 0 && BV % (2 * kU) == 0, "Invalid head dimensions"); // KV seq frags per warp constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - using otile_t = NAXTile; + using otile_t = NAXTile; otile_t Otile; Otile.clear(); @@ -421,8 +424,8 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id += 2) { - if constexpr (BD == 128) { + for (short id = 0; id < TV; id += 2) { + if constexpr (BV == 128) { if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } 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..7e7aba37ae 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 @@ -11,6 +11,12 @@ "_wm" #wm "_wn" #wn "_mask" #mname, \ attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float) +#define instantiate_attn_asymmetric(tname, dtype, bq, bk, bd, bv, wm, wn, mname, mtype) \ + instantiate_kernel( \ + "steel_attention_" #tname "_bq" #bq "_bk" #bk "_bd" #bd "_bv" #bv \ + "_wm" #wm "_wn" #wn "_mask" #mname, \ + attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float, bv) + #define instantiate_attn_dsplit(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ instantiate_kernel( \ "steel_attention_dsplit_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ @@ -27,7 +33,9 @@ #define instantiate_attn_mask_helper(iname, itype) \ instantiate_attn_shapes_helper(iname, itype, iname, itype) \ - instantiate_attn_shapes_helper(iname, itype, bool_, bool) + instantiate_attn_shapes_helper(iname, itype, bool_, bool) \ + instantiate_attn_asymmetric(iname, itype, 64, 32, 96, 64, 4, 1, iname, itype) \ + instantiate_attn_asymmetric(iname, itype, 64, 32, 96, 64, 4, 1, bool_, bool) instantiate_attn_mask_helper(float16, half); instantiate_attn_mask_helper(bfloat16, bfloat); diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index f1141e9792..ba0da0be99 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -505,6 +505,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int, int, int, + int, const array&, bool) { return d.get_kernel(kernel_name, hash_name, func_consts); diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 6ece8c43fb..237eb55c07 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -29,6 +29,7 @@ void sdpa_full_self_attention_nax( using namespace mlx::steel; int bd = q.shape(-1); + int bv = v.shape(-1); int bq = 64; int bk = 32; @@ -98,6 +99,7 @@ void sdpa_full_self_attention_nax( bk, "_bd", bd, + (bv == bd ? "" : "_bv" + std::to_string(bv)), "_wm", wm, "_wn", @@ -131,6 +133,7 @@ void sdpa_full_self_attention_nax( bq, bk, bd, + bv, wm, wn, (has_mask ? *mask : q), @@ -733,17 +736,30 @@ std::tuple has_fused_kernel( std::ostringstream msg; if (query_sequence_length > 8) { - 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); + const bool asymmetric = query_head_dim == 96 && value_head_dim == 64; + const bool supported_head_dim = asymmetric || + (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)); 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 << "."; + << "dims, or (query, value) head dims (96, 64) with NAX; got query " + << "head dim " << query_head_dim << " and value head dim " + << value_head_dim << "."; return {false, msg.str()}; } + if (asymmetric) { + if (!metal::is_nax_available() || + (q.dtype() == float32 && !env::enable_tf32())) { + return { + false, + "the (96, 64) full attention kernel requires NAX and " + "TF32 for float32 inputs."}; + } + } if (has_mask && !has_arr_mask && !(query_sequence_length <= key_sequence_length && do_causal)) { msg << "the full attention kernel with a causal mask requires the " @@ -758,11 +774,12 @@ std::tuple has_fused_kernel( (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || query_head_dim == 192 || query_head_dim == 256 || query_head_dim == 512)) || - (query_head_dim == 192 && value_head_dim == 128); + (query_head_dim == 192 && value_head_dim == 128) || + (query_head_dim == 96 && value_head_dim == 64); if (!supported_head_dim) { msg << "the vector attention kernel supports head dims " << "{64, 96, 128, 192, 256, 512} with matching query/value head " - << "dims, or query head dim 192 with value head dim 128; got " + << "dims, or (query, value) head dims (96, 64) or (192, 128); got " << "query head dim " << query_head_dim << " and value head dim " << value_head_dim << "."; return {false, msg.str()}; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index fefe7325b3..71f9dc2352 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -1,5 +1,6 @@ import math import os +import platform import unittest from itertools import product @@ -292,6 +293,92 @@ 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_d96_v64(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + arch = mx.device_info()["architecture"] + min_gen = 18 if arch.endswith("p") else 17 + if int(arch[-3:-1]) < min_gen or tuple( + map(int, platform.mac_ver()[0].split(".")[:2]) + ) < (26, 2): + self.skipTest("requires NAX") + mx.random.seed(0) + scale = 96**-0.5 + cells = [ + (1, 40, 40, 9, 33), + (1, 4, 4, 64, 64), + (2, 4, 2, 65, 97), + (1, 4, 4, 16, 9), + ] + for dtype, (B, qH, kH, qL, kL) in product( + (mx.float16, mx.bfloat16, mx.float32), cells + ): + if dtype == mx.float32 and os.environ.get("MLX_ENABLE_TF32", "1") == "0": + continue + q = ( + (0.5 * mx.random.normal((B, qL, qH, 96))) + .astype(dtype) + .transpose(0, 2, 1, 3) + ) + k = (0.5 * mx.random.normal((B, kH, kL + 32, 96))).astype(dtype)[:, :, :kL] + masks = [ + None, + mx.arange(kL) < kL - 3, + mx.random.normal((qL, kL)).astype(dtype), + ] + if qL <= kL: + masks.append("causal") + for dv, mask, with_sinks in product((64, 96), masks, (False, True)): + with self.subTest( + dtype=dtype, + cell=(B, qH, kH, qL, kL), + dv=dv, + mask=type(mask).__name__, + sinks=with_sinks, + ): + v = (0.5 * mx.random.normal((B, kH, kL + 32, dv))).astype(dtype)[ + :, :, :kL + ] + if qL == 9: + # V sliced from the interleaved K/V projection. + v = (0.5 * mx.random.normal((B, kL, kH, 2 * dv))).astype(dtype) + v = v.transpose(0, 2, 1, 3)[..., dv:] + sinks = ( + mx.random.normal((qH,)).astype(dtype) if with_sinks else None + ) + ref = mlx_ref_attn( + q.astype(mx.float32), + k.astype(mx.float32), + v.astype(mx.float32), + scale, + mask, + sinks, + ) + out = mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=scale, + mask=mask, + sinks=sinks, + force_fused=True, + ) + self.assertEqual(out.shape, (B, qH, qL, dv)) + tol = 5e-3 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_full_d96_rejects_unsupported_value_dim(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + q = mx.random.normal((1, 4, 9, 96)).astype(mx.float16) + k = mx.random.normal((1, 4, 33, 96)).astype(mx.float16) + v = mx.random.normal((1, 4, 33, 80)).astype(mx.float16) + scale = 96**-0.5 + with self.assertRaisesRegex(ValueError, "supports head dims"): + mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, force_fused=True) + def test_sdpa_vector_kv_transposed_head_seq(self): D = 64 Nq = 4 @@ -410,6 +497,51 @@ def test_sdpa_vector_gqa_long(self): out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_vector_d96_v64(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + mx.random.seed(0) + scale = 96**-0.5 + # Include MiniCPM3 heads and GQA to reach two-pass on smaller GPUs. + cells = [(1, 40, 40, 1, 127), (1, 40, 40, 8, 1025), (2, 4, 2, 4, 4097)] + for dtype, (B, qH, kH, qL, kL) in product( + (mx.float32, mx.float16, mx.bfloat16), cells + ): + q = mx.random.normal((B, qL, qH, 96)).astype(dtype) + q = q.transpose(0, 2, 1, 3) + k = mx.random.normal((B, kH, kL + 32, 96)).astype(dtype)[:, :, :kL] + v = mx.random.normal((B, kH, kL + 32, 64)).astype(dtype)[:, :, :kL] + masks = ( + None, + "causal", + mx.arange(kL) < kL - 3, + mx.random.normal((qL, kL)).astype(dtype), + ) + for mask in masks: + with self.subTest( + dtype=dtype, cell=(B, qH, kH, qL, kL), mask=type(mask).__name__ + ): + # GPU float32 GEMM may use TF32; keep the reference exact. + with mx.stream(mx.cpu): + ref = mlx_ref_attn( + q.astype(mx.float32), + k.astype(mx.float32), + v.astype(mx.float32), + scale, + mask, + ) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask, force_fused=True + ) + self.assertEqual(out.shape, (B, qH, qL, 64)) + tol = ( + 1e-5 + if dtype == mx.float32 + else (3e-3 if dtype == mx.float16 else 2e-2) + ) + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") def test_sdpa_vector_head_dim_512(self): if mx.default_device() != mx.gpu: