Run the vector SDPA kernels at head_dim 256 on CUDA

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<int, 128>{});
       break;
+    case 256:
+      f(std::integral_constant<int, 256>{});
+      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;

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