diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cu b/mlx/backend/cuda/scaled_dot_product_attention.cu index 6fafb0ba98..c6996c95ef 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cu +++ b/mlx/backend/cuda/scaled_dot_product_attention.cu @@ -460,6 +460,9 @@ void dispatch_headdim(int n, F&& f) { case 128: f(std::integral_constant{}); break; + case 256: + f(std::integral_constant{}); + break; } } @@ -674,7 +677,8 @@ bool supports_sdpa_vector( const int key_sequence_length = k.shape(2); const bool sdpa_supported_head_dim = query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128); + (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || + query_head_dim == 256); const bool supported_vector_config = sdpa_supported_head_dim && query_sequence_length < 4; diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index b2d7e9121e..b47ed1bdae 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -622,7 +622,13 @@ def test_sdpa_vector(self): def test_sdpa_vector_gqa_long(self): scale = 1.0 mx.random.seed(0) - for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64), (48, 4, 128), (64, 4, 128)]: + for Nq, Nkv, D in [ + (32, 4, 128), + (64, 8, 64), + (48, 4, 128), + (64, 4, 128), + (16, 2, 256), + ]: for B, L in [(1, 8192), (1, 8201), (2, 8192)]: q = 5e-1 * mx.random.normal(shape=(B, Nq, 1, D)) k = 5e-1 * mx.random.normal(shape=(B, Nkv, L + 32, D))[:, :, :L] @@ -1350,7 +1356,7 @@ def make_qkv(qL, kL, D, qH=8, kH=8): return q, k, v # Vector attention kernel. - for D in (64, 96, 128): + for D in (64, 96, 128, 256): with self.subTest(head_dim=D): q, k, v = make_qkv(3, 128, D, 4, 2) scale = D**-0.5