Fix

Author: mseeger

tests/test_model.py

@@ -35,6 +35,7 @@
 
 import litgpt.config as config_module
 from litgpt import GPT, Config
+from litgpt.attention import DefaultKeysAndValues
 from litgpt.model import CausalSelfAttention
 from litgpt.scripts.convert_hf_checkpoint import (
     copy_weights_falcon,
@@ -1495,6 +1496,18 @@ 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)
+            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)