From 7ad38d432357e3d3766763fe0771fdd0aa60e1f7 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Wed, 2 Sep 2026 01:44:28 -0700 Subject: [PATCH 1/8] Enable fused vector attention for head dimension 512 Instantiate the single- and two-pass Metal vector attention kernels for matching query and value dimensions of 512, and admit that dimension in the short-query fused dispatch. Add coverage for single- and two-pass lengths, float32/float16/bfloat16, query lengths 1 and 4, boolean and additive masks, non-contiguous masks, causal masking, batches, attention sinks, and force_fused engagement. --- .../scaled_dot_product_attention.metal | 4 +- .../metal/scaled_dot_product_attention.cpp | 8 +-- python/tests/test_fast_sdpa.py | 67 ++++++++++++++++++- 3 files changed, 73 insertions(+), 6 deletions(-) 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..c0dda466c3 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()}; } diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 6b185ecf1f..27a496120b 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -418,6 +418,71 @@ 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.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_vector_head_dim_512(self): + # 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) + + # L = 128 takes the 1-pass kernel, 8192 and 8201 the 2-pass one. + for L, dtype in product( + (128, 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) + atol = 1e-5 if dtype == mx.float32 else 2e-2 + self.assertTrue(mx.allclose(ref, out, atol=atol)) + + # The vector kernel fuses while query length * GQA factor is at most + # 32, so with 8 repeats the query length can be up to 4. + L = 1024 + k = mx.random.normal(shape=(1, Nkv, L, D)) + v = mx.random.normal(shape=(1, Nkv, L, D)) + for qL in (1, 4): + q = mx.random.normal(shape=(1, Nq, qL, D)) + masks = [ + None, + mx.array(True), + mx.array([True] * (L - 10) + [False] * 10), + mx.random.uniform(shape=(Nq, qL, L)) > 0.2, + mx.random.uniform(shape=(L, qL, Nq)).T > 0.2, + mx.random.uniform(shape=(Nq, qL, L)), + mx.random.uniform(shape=(L, qL, Nq)).T, + "causal", + ] + # The reference does the multi-query scores as a GEMM, which is + # the less accurate arm at qL > 1, so use the same tolerance as + # the other multi-query tests. + tol = 1e-4 if qL == 1 else 1e-3 + for i, m in enumerate(masks): + with self.subTest(qL=qL, mask=i): + ref = mlx_ref_attn(q, k, v, scale, mask=m) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=m + ) + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + + # Batched, with and without attention sinks, over both kernels. + 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 + ) + 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,7 +953,7 @@ 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 From c65be762fc498781d1fe4290564003b647fa0e33 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Fri, 4 Sep 2026 01:03:39 -0700 Subject: [PATCH 2/8] Gate head dim 512 vector admission on a minimum key length The D=512 cells run the generic (non-GQA) vector kernel, whose cost is dominated by reading K/V: the fused win over the unfused path grows with the key length and is only a few percent at short ones. Admit head dim 512 only at or above MLX_SDPA_D512_MIN_KL keys (default 1024, read per call) and keep the existing fallback below that. Cover the threshold in both tests: force_fused at a short key length raises with the new refusal, and the threshold override restores the 1-pass kernel at key lengths the default excludes. --- .../metal/scaled_dot_product_attention.cpp | 16 +++++ python/tests/test_fast_sdpa.py | 59 ++++++++++++++++--- 2 files changed, 68 insertions(+), 7 deletions(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index c0dda466c3..3ee9ab7176 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -780,6 +780,22 @@ std::tuple has_fused_kernel( << query_sequence_length << " and GQA factor " << gqa_factor << "."; return {false, msg.str()}; } + // Head dim 512 sits on the generic (non-GQA) vector kernel, whose cost + // is dominated by reading K/V. Its measured win over the unfused path + // grows with the key length and is only a few percent at short ones, + // so admit it from a minimum key length and keep the existing path + // below that. The threshold is read per call so callers and tests can + // observe both routings in one process. + if (query_head_dim == 512) { + const int min_key_sequence_length = + env::get_var("MLX_SDPA_D512_MIN_KL", 1024); + 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 27a496120b..ac2c3615ab 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -426,9 +426,39 @@ def test_sdpa_vector_head_dim_512(self): scale = D**-0.5 mx.random.seed(0) - # L = 128 takes the 1-pass kernel, 8192 and 8201 the 2-pass one. + # Below the admission threshold the generic vector kernel stays on + # the unfused path, and force_fused says so. The threshold is read + # per call from MLX_SDPA_D512_MIN_KL (default 1024). + q = mx.random.normal(shape=(1, Nq, 1, D), dtype=mx.float16) + k = mx.random.normal(shape=(1, Nkv, 128, D), dtype=mx.float16) + v = mx.random.normal(shape=(1, Nkv, 128, D), dtype=mx.float16) + with self.assertRaisesRegex( + ValueError, r"requires at least \d+ keys" + ): + mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + + # L = 128 admitted via the threshold override takes the 1-pass + # kernel; 8192 and 8201 take the 2-pass one at the default + # threshold. + os.environ["MLX_SDPA_D512_MIN_KL"] = "0" + try: + 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 + ) + atol = 1e-5 if dtype == mx.float32 else 2e-2 + self.assertTrue(mx.allclose(ref, out, atol=atol)) + finally: + os.environ.pop("MLX_SDPA_D512_MIN_KL", None) for L, dtype in product( - (128, 8192, 8201), (mx.float32, mx.float16, mx.bfloat16) + (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) @@ -469,6 +499,8 @@ def test_sdpa_vector_head_dim_512(self): self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) # Batched, with and without attention sinks, over both kernels. + # L = 256 is below the admission threshold, so the fused-1-pass leg + # runs with the threshold override; L = 8192 fuses by default. B = 2 sinks = 10 * mx.random.normal(shape=(Nq,)) q = mx.random.normal(shape=(B, Nq, 1, D)) @@ -477,11 +509,19 @@ def test_sdpa_vector_head_dim_512(self): 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 - ) - self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + if L == 256: + os.environ["MLX_SDPA_D512_MIN_KL"] = "0" + try: + 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 + ) + self.assertTrue( + mx.allclose(ref, out, atol=1e-4, rtol=1e-4) + ) + finally: + if L == 256: + os.environ.pop("MLX_SDPA_D512_MIN_KL", None) def test_sdpa_fully_masked(self): Lkv = 8 @@ -993,6 +1033,11 @@ 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"): + 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): From 0ca6a07f0a9e932abc7adb2044c32e642790892a Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Fri, 4 Sep 2026 03:21:37 -0700 Subject: [PATCH 3/8] Harden D512 threshold coverage tests --- python/tests/test_fast_sdpa.py | 41 ++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 14 deletions(-) diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index ac2c3615ab..c372605fcf 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 @@ -426,6 +428,13 @@ def test_sdpa_vector_head_dim_512(self): 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.pop("MLX_SDPA_D512_MIN_KL", None) + # Below the admission threshold the generic vector kernel stays on # the unfused path, and force_fused says so. The threshold is read # per call from MLX_SDPA_D512_MIN_KL (default 1024). @@ -442,8 +451,7 @@ def test_sdpa_vector_head_dim_512(self): # L = 128 admitted via the threshold override takes the 1-pass # kernel; 8192 and 8201 take the 2-pass one at the default # threshold. - os.environ["MLX_SDPA_D512_MIN_KL"] = "0" - try: + with patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}): 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) @@ -451,12 +459,10 @@ def test_sdpa_vector_head_dim_512(self): 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 + 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)) - finally: - os.environ.pop("MLX_SDPA_D512_MIN_KL", None) for L, dtype in product( (8192, 8201), (mx.float32, mx.float16, mx.bfloat16) ): @@ -465,7 +471,9 @@ def test_sdpa_vector_head_dim_512(self): 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) + 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)) @@ -494,7 +502,7 @@ def test_sdpa_vector_head_dim_512(self): with self.subTest(qL=qL, mask=i): ref = mlx_ref_attn(q, k, v, scale, mask=m) out = mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, mask=m + q, k, v, scale=scale, mask=m, force_fused=True ) self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) @@ -509,19 +517,24 @@ def test_sdpa_vector_head_dim_512(self): 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): - if L == 256: - os.environ["MLX_SDPA_D512_MIN_KL"] = "0" - try: + threshold = ( + patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}) + if L == 256 + else nullcontext() + ) + with threshold: 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 + q, + k, + v, + scale=scale, + sinks=s_in, + force_fused=True, ) self.assertTrue( mx.allclose(ref, out, atol=1e-4, rtol=1e-4) ) - finally: - if L == 256: - os.environ.pop("MLX_SDPA_D512_MIN_KL", None) def test_sdpa_fully_masked(self): Lkv = 8 From 2413e44444b6f1348960277738cbc69ca733b2e3 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Fri, 4 Sep 2026 03:36:31 -0700 Subject: [PATCH 4/8] Narrow D512 vector attention admission --- .../metal/scaled_dot_product_attention.cpp | 29 +++++-- python/tests/test_fast_sdpa.py | 77 ++++++++++++++++--- 2 files changed, 90 insertions(+), 16 deletions(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index 3ee9ab7176..c0eaf8e11b 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -780,15 +780,32 @@ std::tuple has_fused_kernel( << query_sequence_length << " and GQA factor " << gqa_factor << "."; return {false, msg.str()}; } - // Head dim 512 sits on the generic (non-GQA) vector kernel, whose cost - // is dominated by reading K/V. Its measured win over the unfused path - // grows with the key length and is only a few percent at short ones, - // so admit it from a minimum key length and keep the existing path - // below that. The threshold is read per call so callers and tests can - // observe both routings in one process. + // Head dim 512 sits on the generic vector kernel, whose cost is dominated + // by reading K/V. Default admission is limited to the measured decode + // surface: a single query, GQA factor 8, and no array mask. + // The threshold is read per call; setting it to zero deliberately expands + // admission for benchmarking and correctness coverage. if (query_head_dim == 512) { const int min_key_sequence_length = env::get_var("MLX_SDPA_D512_MIN_KL", 1024); + const bool expand_admission = min_key_sequence_length == 0; + if (!expand_admission && 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 (!expand_admission && 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 (!expand_admission && 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 " diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index c372605fcf..5711c41afe 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -477,11 +477,53 @@ def test_sdpa_vector_head_dim_512(self): atol = 1e-5 if dtype == mx.float32 else 2e-2 self.assertTrue(mx.allclose(ref, out, atol=atol)) - # The vector kernel fuses while query length * GQA factor is at most - # 32, so with 8 repeats the query length can be up to 4. + # By default D=512 admits the measured single-token GQA surface, not + # MHA, multi-token queries, or array masks. The zero threshold is an + # explicit expansion used to retain correctness coverage for those + # otherwise supported vector-kernel shapes. L = 1024 k = mx.random.normal(shape=(1, Nkv, L, D)) v = mx.random.normal(shape=(1, Nkv, L, D)) + + q = mx.random.normal(shape=(1, Nq, 4, D)) + with self.assertRaisesRegex(ValueError, r"single-token queries"): + mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + + q = mx.random.normal(shape=(1, Nq, 1, D)) + array_mask = mx.array([True] * (L - 10) + [False] * 10) + with self.assertRaisesRegex(ValueError, r"without array masks"): + mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=array_mask, force_fused=True + ) + + for q_heads, kv_heads in ((8, 8), (32, 8)): + with self.subTest(q_heads=q_heads, kv_heads=kv_heads): + other_q = mx.random.normal(shape=(1, q_heads, 1, D)) + other_k = mx.random.normal(shape=(1, kv_heads, L, D)) + other_v = mx.random.normal(shape=(1, kv_heads, L, D)) + with self.assertRaisesRegex(ValueError, r"GQA factor 8"): + mx.fast.scaled_dot_product_attention( + other_q, + other_k, + other_v, + scale=scale, + force_fused=True, + ) + with patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}): + ref = mlx_ref_attn(other_q, other_k, other_v, scale) + out = mx.fast.scaled_dot_product_attention( + other_q, + other_k, + other_v, + scale=scale, + force_fused=True, + ) + self.assertTrue( + mx.allclose(ref, out, atol=1e-4, rtol=1e-4) + ) + for qL in (1, 4): q = mx.random.normal(shape=(1, Nq, qL, D)) masks = [ @@ -500,11 +542,20 @@ def test_sdpa_vector_head_dim_512(self): tol = 1e-4 if qL == 1 else 1e-3 for i, m in enumerate(masks): with self.subTest(qL=qL, mask=i): - ref = mlx_ref_attn(q, k, v, scale, mask=m) - out = mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, mask=m, force_fused=True + needs_override = qL > 1 or i not in (0, 7) + threshold = ( + patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}) + if needs_override + else nullcontext() ) - self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + with threshold: + ref = mlx_ref_attn(q, k, v, scale, mask=m) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=m, force_fused=True + ) + self.assertTrue( + mx.allclose(ref, out, atol=tol, rtol=tol) + ) # Batched, with and without attention sinks, over both kernels. # L = 256 is below the admission threshold, so the fused-1-pass leg @@ -1010,11 +1061,17 @@ def make_qkv(qL, kL, D, qH=8, kH=8): 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"): From b486f2ec8c0ec438fc0d1e62e2afb001c60f4908 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Sat, 5 Sep 2026 15:51:54 -0700 Subject: [PATCH 5/8] Isolate the default D512 threshold assertion from the environment --- python/tests/test_fast_sdpa.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 5711c41afe..b492b9c971 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -1103,11 +1103,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"): - 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 - ) + with patch.dict(os.environ): + os.environ.pop("MLX_SDPA_D512_MIN_KL", None) + with self.assertRaisesRegex(ValueError, r"requires at least \d+ keys"): + 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): From d2f46473f195444a98c282c37091bf734350512d Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Sun, 6 Sep 2026 17:07:59 -0700 Subject: [PATCH 6/8] Limit D512 Metal routing assertions to Metal GPU tests --- python/tests/test_fast_sdpa.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index b492b9c971..befadcf405 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -420,8 +420,10 @@ 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.is_available(mx.gpu), "GPU kernel path only") + @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 From c30bc7bb1d9a6300e06a595095874a6c4764c603 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Sun, 6 Sep 2026 20:38:16 -0700 Subject: [PATCH 7/8] Apply pre-commit formatting --- python/tests/test_fast_sdpa.py | 28 +++++++--------------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index befadcf405..d4167f5b78 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -443,12 +443,8 @@ def test_sdpa_vector_head_dim_512(self): q = mx.random.normal(shape=(1, Nq, 1, D), dtype=mx.float16) k = mx.random.normal(shape=(1, Nkv, 128, D), dtype=mx.float16) v = mx.random.normal(shape=(1, Nkv, 128, D), dtype=mx.float16) - with self.assertRaisesRegex( - ValueError, r"requires at least \d+ keys" - ): - mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, force_fused=True - ) + with self.assertRaisesRegex(ValueError, r"requires at least \d+ keys"): + mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, force_fused=True) # L = 128 admitted via the threshold override takes the 1-pass # kernel; 8192 and 8201 take the 2-pass one at the default @@ -465,9 +461,7 @@ def test_sdpa_vector_head_dim_512(self): ) atol = 1e-5 if dtype == mx.float32 else 2e-2 self.assertTrue(mx.allclose(ref, out, atol=atol)) - for L, dtype in product( - (8192, 8201), (mx.float32, mx.float16, mx.bfloat16) - ): + 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) @@ -489,9 +483,7 @@ def test_sdpa_vector_head_dim_512(self): q = mx.random.normal(shape=(1, Nq, 4, D)) with self.assertRaisesRegex(ValueError, r"single-token queries"): - mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, force_fused=True - ) + mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, force_fused=True) q = mx.random.normal(shape=(1, Nq, 1, D)) array_mask = mx.array([True] * (L - 10) + [False] * 10) @@ -522,9 +514,7 @@ def test_sdpa_vector_head_dim_512(self): scale=scale, force_fused=True, ) - self.assertTrue( - mx.allclose(ref, out, atol=1e-4, rtol=1e-4) - ) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) for qL in (1, 4): q = mx.random.normal(shape=(1, Nq, qL, D)) @@ -555,9 +545,7 @@ def test_sdpa_vector_head_dim_512(self): out = mx.fast.scaled_dot_product_attention( q, k, v, scale=scale, mask=m, force_fused=True ) - self.assertTrue( - mx.allclose(ref, out, atol=tol, rtol=tol) - ) + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) # Batched, with and without attention sinks, over both kernels. # L = 256 is below the admission threshold, so the fused-1-pass leg @@ -585,9 +573,7 @@ def test_sdpa_vector_head_dim_512(self): sinks=s_in, force_fused=True, ) - self.assertTrue( - mx.allclose(ref, out, atol=1e-4, rtol=1e-4) - ) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) def test_sdpa_fully_masked(self): Lkv = 8 From 81d837b7a4f924b19dafd7bfb194308e0d991add Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 8 Sep 2026 16:21:35 +0900 Subject: [PATCH 8/8] nit --- .../metal/scaled_dot_product_attention.cpp | 19 +-- python/tests/test_fast_sdpa.py | 160 +++++------------- 2 files changed, 50 insertions(+), 129 deletions(-) diff --git a/mlx/backend/metal/scaled_dot_product_attention.cpp b/mlx/backend/metal/scaled_dot_product_attention.cpp index c0eaf8e11b..6ece8c43fb 100644 --- a/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -780,28 +780,25 @@ std::tuple has_fused_kernel( << query_sequence_length << " and GQA factor " << gqa_factor << "."; return {false, msg.str()}; } - // Head dim 512 sits on the generic vector kernel, whose cost is dominated - // by reading K/V. Default admission is limited to the measured decode - // surface: a single query, GQA factor 8, and no array mask. - // The threshold is read per call; setting it to zero deliberately expands - // admission for benchmarking and correctness coverage. + // 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) { - const int min_key_sequence_length = - env::get_var("MLX_SDPA_D512_MIN_KL", 1024); - const bool expand_admission = min_key_sequence_length == 0; - if (!expand_admission && query_sequence_length != 1) { + 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 (!expand_admission && gqa_factor != 8) { + 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 (!expand_admission && has_arr_mask) { + if (!always && has_arr_mask) { msg << "the vector attention kernel for head dim 512 defaults to " "calls without array masks."; return {false, msg.str()}; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index d4167f5b78..48f1527a66 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -435,32 +435,22 @@ def test_sdpa_vector_head_dim_512(self): environment = patch.dict(os.environ) environment.start() self.addCleanup(environment.stop) - os.environ.pop("MLX_SDPA_D512_MIN_KL", None) - - # Below the admission threshold the generic vector kernel stays on - # the unfused path, and force_fused says so. The threshold is read - # per call from MLX_SDPA_D512_MIN_KL (default 1024). - q = mx.random.normal(shape=(1, Nq, 1, D), dtype=mx.float16) - k = mx.random.normal(shape=(1, Nkv, 128, D), dtype=mx.float16) - v = mx.random.normal(shape=(1, Nkv, 128, D), dtype=mx.float16) - with self.assertRaisesRegex(ValueError, r"requires at least \d+ keys"): - mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, force_fused=True) - - # L = 128 admitted via the threshold override takes the 1-pass - # kernel; 8192 and 8201 take the 2-pass one at the default - # threshold. - with patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}): - 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)) + 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) @@ -473,83 +463,23 @@ def test_sdpa_vector_head_dim_512(self): atol = 1e-5 if dtype == mx.float32 else 2e-2 self.assertTrue(mx.allclose(ref, out, atol=atol)) - # By default D=512 admits the measured single-token GQA surface, not - # MHA, multi-token queries, or array masks. The zero threshold is an - # explicit expansion used to retain correctness coverage for those - # otherwise supported vector-kernel shapes. - L = 1024 - k = mx.random.normal(shape=(1, Nkv, L, D)) - v = mx.random.normal(shape=(1, Nkv, L, D)) - - q = mx.random.normal(shape=(1, Nq, 4, D)) - with self.assertRaisesRegex(ValueError, r"single-token queries"): - mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, force_fused=True) - - q = mx.random.normal(shape=(1, Nq, 1, D)) - array_mask = mx.array([True] * (L - 10) + [False] * 10) - with self.assertRaisesRegex(ValueError, r"without array masks"): - mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, mask=array_mask, force_fused=True - ) - + # Test other heads. for q_heads, kv_heads in ((8, 8), (32, 8)): with self.subTest(q_heads=q_heads, kv_heads=kv_heads): - other_q = mx.random.normal(shape=(1, q_heads, 1, D)) - other_k = mx.random.normal(shape=(1, kv_heads, L, D)) - other_v = mx.random.normal(shape=(1, kv_heads, L, D)) - with self.assertRaisesRegex(ValueError, r"GQA factor 8"): - mx.fast.scaled_dot_product_attention( - other_q, - other_k, - other_v, - scale=scale, - force_fused=True, - ) - with patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}): - ref = mlx_ref_attn(other_q, other_k, other_v, scale) - out = mx.fast.scaled_dot_product_attention( - other_q, - other_k, - other_v, - scale=scale, - force_fused=True, - ) - self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) - - for qL in (1, 4): - q = mx.random.normal(shape=(1, Nq, qL, D)) - masks = [ - None, - mx.array(True), - mx.array([True] * (L - 10) + [False] * 10), - mx.random.uniform(shape=(Nq, qL, L)) > 0.2, - mx.random.uniform(shape=(L, qL, Nq)).T > 0.2, - mx.random.uniform(shape=(Nq, qL, L)), - mx.random.uniform(shape=(L, qL, Nq)).T, - "causal", - ] - # The reference does the multi-query scores as a GEMM, which is - # the less accurate arm at qL > 1, so use the same tolerance as - # the other multi-query tests. - tol = 1e-4 if qL == 1 else 1e-3 - for i, m in enumerate(masks): - with self.subTest(qL=qL, mask=i): - needs_override = qL > 1 or i not in (0, 7) - threshold = ( - patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}) - if needs_override - else nullcontext() - ) - with threshold: - ref = mlx_ref_attn(q, k, v, scale, mask=m) - out = mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, mask=m, force_fused=True - ) - self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + 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)) - # Batched, with and without attention sinks, over both kernels. - # L = 256 is below the admission threshold, so the fused-1-pass leg - # runs with the threshold override; L = 8192 fuses by default. + # Test batched. B = 2 sinks = 10 * mx.random.normal(shape=(Nq,)) q = mx.random.normal(shape=(B, Nq, 1, D)) @@ -558,22 +488,16 @@ def test_sdpa_vector_head_dim_512(self): 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): - threshold = ( - patch.dict(os.environ, {"MLX_SDPA_D512_MIN_KL": "0"}) - if L == 256 - else nullcontext() + 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, ) - with threshold: - 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)) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) def test_sdpa_fully_masked(self): Lkv = 8 @@ -1091,9 +1015,9 @@ 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 patch.dict(os.environ): - os.environ.pop("MLX_SDPA_D512_MIN_KL", None) - with self.assertRaisesRegex(ValueError, r"requires at least \d+ keys"): + 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