diff --git a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index db86a0edea..68e2a6c06a 100644 --- a/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/mlx/backend/metal/kernels/scaled_dot_product_attention.metal @@ -45,6 +45,7 @@ using namespace metal; instantiate_sdpa_vector(type, 192, 128) \ instantiate_sdpa_vector(type, 192, 192) \ instantiate_sdpa_vector(type, 256, 256) \ + instantiate_sdpa_vector(type, 512, 512) \ instantiate_sdpa_vector_gqa(type, 64, 64, 8, 8) \ instantiate_sdpa_vector_gqa(type, 128, 128, 8, 4) \ instantiate_sdpa_vector_gqa(type, 128, 128, 12, 4) \ @@ -53,7 +54,8 @@ using namespace metal; instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ instantiate_sdpa_vector_aggregation(type, 192) \ - instantiate_sdpa_vector_aggregation(type, 256) + instantiate_sdpa_vector_aggregation(type, 256) \ + instantiate_sdpa_vector_aggregation(type, 512) instantiate_sdpa_vector_heads(float) instantiate_sdpa_vector_heads(bfloat16_t) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 6d9ade1ccc..6ece8c43fb 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -757,13 +757,13 @@ std::tuple has_fused_kernel( (query_head_dim == value_head_dim && (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || query_head_dim == 192 || - query_head_dim == 256)) || + query_head_dim == 256 || query_head_dim == 512)) || (query_head_dim == 192 && value_head_dim == 128); if (!supported_head_dim) { msg << "the vector attention kernel supports head dims " - << "{64, 96, 128, 192, 256} with matching query/value head dims, " - << "or query head dim 192 with value head dim 128; got query head " - << "dim " << query_head_dim << " and value head dim " + << "{64, 96, 128, 192, 256, 512} with matching query/value head " + << "dims, or query head dim 192 with value head dim 128; got " + << "query head dim " << query_head_dim << " and value head dim " << value_head_dim << "."; return {false, msg.str()}; } @@ -780,6 +780,36 @@ std::tuple has_fused_kernel( << query_sequence_length << " and GQA factor " << gqa_factor << "."; return {false, msg.str()}; } + // Head dim 512 uses the generic vector kernel, reading K/V is its main + // cost. By default, use it only for one query, GQA factor 8, and no array + // mask. + if (query_head_dim == 512) { + int min_key_sequence_length = env::get_var("MLX_SDPA_D512_MIN_KL", 1024); + bool always = (min_key_sequence_length == 0); + if (!always && query_sequence_length != 1) { + msg << "the vector attention kernel for head dim 512 defaults to " + "single-token queries; got query length " + << query_sequence_length << "."; + return {false, msg.str()}; + } + if (!always && gqa_factor != 8) { + msg << "the vector attention kernel for head dim 512 defaults to " + "GQA factor 8; got GQA factor " + << gqa_factor << "."; + return {false, msg.str()}; + } + if (!always && has_arr_mask) { + msg << "the vector attention kernel for head dim 512 defaults to " + "calls without array masks."; + return {false, msg.str()}; + } + if (key_sequence_length < min_key_sequence_length) { + msg << "the vector attention kernel for head dim 512 requires at " + << "least " << min_key_sequence_length << " keys; got key length " + << key_sequence_length << "."; + return {false, msg.str()}; + } + } } return {true, ""}; } diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 6b185ecf1f..48f1527a66 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -1,7 +1,9 @@ import math import os import unittest +from contextlib import nullcontext from itertools import product +from unittest.mock import patch import mlx.core as mx import mlx_tests @@ -418,6 +420,85 @@ 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_head_dim_512(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + # gemma-4 global attention: 32 query heads over 4 key/value heads. + D = 512 + Nq, Nkv = 32, 4 + scale = D**-0.5 + mx.random.seed(0) + + # Keep the test hermetic if the caller already configured a threshold. + # addCleanup restores the complete environment snapshot on every exit. + environment = patch.dict(os.environ) + environment.start() + self.addCleanup(environment.stop) + os.environ["MLX_SDPA_D512_MIN_KL"] = "0" + + # Test 1-pass kernel. + for dtype in (mx.float32, mx.float16, mx.bfloat16): + with self.subTest(L=128, dtype=dtype, threshold="off"): + q = mx.random.normal(shape=(1, Nq, 1, D), dtype=dtype) + k = mx.random.normal(shape=(1, Nkv, 128, D), dtype=dtype) + v = mx.random.normal(shape=(1, Nkv, 128, D), dtype=dtype) + ref = mlx_ref_attn(q, k, v, scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + atol = 1e-5 if dtype == mx.float32 else 2e-2 + self.assertTrue(mx.allclose(ref, out, atol=atol)) + + # Test 2-pass kernel. + for L, dtype in product((8192, 8201), (mx.float32, mx.float16, mx.bfloat16)): + with self.subTest(L=L, dtype=dtype): + q = mx.random.normal(shape=(1, Nq, 1, D), dtype=dtype) + k = mx.random.normal(shape=(1, Nkv, L, D), dtype=dtype) + v = mx.random.normal(shape=(1, Nkv, L, D), dtype=dtype) + ref = mlx_ref_attn(q, k, v, scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + atol = 1e-5 if dtype == mx.float32 else 2e-2 + self.assertTrue(mx.allclose(ref, out, atol=atol)) + + # Test other heads. + for q_heads, kv_heads in ((8, 8), (32, 8)): + with self.subTest(q_heads=q_heads, kv_heads=kv_heads): + q = mx.random.normal(shape=(1, q_heads, 1, D)) + k = mx.random.normal(shape=(1, kv_heads, L, D)) + v = mx.random.normal(shape=(1, kv_heads, L, D)) + 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-4, rtol=1e-4)) + + # Test batched. + B = 2 + sinks = 10 * mx.random.normal(shape=(Nq,)) + q = mx.random.normal(shape=(B, Nq, 1, D)) + for L in (256, 8192): + k = mx.random.normal(shape=(B, Nkv, L, D)) + v = mx.random.normal(shape=(B, Nkv, L, D)) + for s_in in (None, sinks): + with self.subTest(L=L, sinks=s_in is not None): + ref = mlx_ref_attn(q, k, v, scale, sinks=s_in) + out = mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=scale, + sinks=s_in, + force_fused=True, + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + def test_sdpa_fully_masked(self): Lkv = 8 mask = mx.array(False) @@ -888,15 +969,21 @@ def make_qkv(qL, kL, D, qH=8, kH=8): self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) # Vector attention kernel. - for D in (192, 256): + for D in (192, 256, 512): with self.subTest(head_dim=D): q, k, v = make_qkv(4, 16385, D, 4, 2) scale = D**-0.5 - ref = mlx_ref_attn(q, k, v, scale=scale) - out = mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, force_fused=True + threshold = ( + patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}) + if D == 512 + else nullcontext() ) - self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + with threshold: + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) # No full attention fused kernels. with self.assertRaisesRegex(ValueError, "supports head dims"): @@ -928,6 +1015,13 @@ def make_qkv(qL, kL, D, qH=8, kH=8): mx.fast.scaled_dot_product_attention( q, k, v, scale=64**-0.5, force_fused=True ) + with self.assertRaisesRegex(ValueError, r"requires at least \d+ keys"): + with patch.dict(os.environ): + os.environ.pop("MLX_SDPA_D512_MIN_KL", None) + q, k, v = make_qkv(1, 512, 512, qH=32, kH=4) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) # No CPU fused kernel. with mx.stream(mx.cpu):