Fix

Author: mseeger

tests/test_model.py

@@ -1444,6 +1444,19 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
         # SDPAParams gained an additional argument in PyTorch 2.5
         args = []
         assert k_and_v.both_in_parallel()
+        # This is also done in `MultiHeadSelfAttention.scaled_dot_product_attention`
+        if mask is None and enable_gqa:
+            # Some efficient kernels have not implemented
+            # `enabla_gqa=True`. It is better to extend keys, values in
+            # this case.
+            key = k_and_v.keys()
+            value = k_and_v.values()
+            q_per_kv = config.n_head // config.n_query_groups
+            key = key.repeat_interleave(q_per_kv, dim=1)
+            value = value.repeat_interleave(q_per_kv, dim=1)
+            assert query.shape[1] == key.shape[1]
+            k_and_v = DefaultKeysAndValues(key, value)
+
         if hasattr(SDPAParams, "enable_gqa"):
             args.append(enable_gqa)
         params = SDPAParams(query, k_and_v.keys(), k_and_v.values(), mask, 0.0, True, *args)
@@ -1506,6 +1519,7 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             q_per_kv = config.n_head // config.n_query_groups
             key = key.repeat_interleave(q_per_kv, dim=1)
             value = value.repeat_interleave(q_per_kv, dim=1)
+            assert query.shape[1] == key.shape[1]
             k_and_v = DefaultKeysAndValues(key, value)
 
         if hasattr(SDPAParams, "enable_gqa"):