Enable fused SDPA vector kernel for asymmetric Q/V head dims (192, 128) (#3637)

Co-authored-by: Cheng <git@zcbenz.com>

Author: yohann-bearzi

mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -32,6 +32,7 @@ using namespace metal;
   instantiate_sdpa_vector(type, 64, 64)          \
   instantiate_sdpa_vector(type, 96, 96)          \
   instantiate_sdpa_vector(type, 128, 128)        \
+  instantiate_sdpa_vector(type, 192, 128)        \
   instantiate_sdpa_vector(type, 256, 256)        \
   instantiate_sdpa_vector_aggregation(type, 64)  \
   instantiate_sdpa_vector_aggregation(type, 96)  \

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -619,9 +619,10 @@ bool ScaledDotProductAttention::use_fallback(
   const int gqa_factor = num_query_heads / num_kv_heads;
 
   const bool sdpa_vector_supported_head_dim =
-      query_head_dim == value_head_dim &&
-      (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 ||
-       query_head_dim == 256);
+      (query_head_dim == value_head_dim &&
+       (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 ||
+        query_head_dim == 256)) ||
+      (query_head_dim == 192 && value_head_dim == 128);
   const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim &&
       (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128);