Fix

Author: mseeger

tests/test_model.py

@@ -1456,9 +1456,12 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             value = value.repeat_interleave(q_per_kv, dim=1)
             assert query.shape[1] == key.shape[1]
             k_and_v = DefaultKeysAndValues(key, value)
+            _enable_gqa = False
+        else:
+            _enable_gqa = enable_gqa
 
         if hasattr(SDPAParams, "enable_gqa"):
-            args.append(enable_gqa)
+            args.append(_enable_gqa)
         params = SDPAParams(query, k_and_v.keys(), k_and_v.values(), mask, 0.0, True, *args)
         if expected is SDPBackend.FLASH_ATTENTION:
             assert flash_sdp_enabled(), "flash_sdp_enabled() is False"
@@ -1521,9 +1524,12 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             value = value.repeat_interleave(q_per_kv, dim=1)
             assert query.shape[1] == key.shape[1]
             k_and_v = DefaultKeysAndValues(key, value)
+            _enable_gqa = False
+        else:
+            _enable_gqa = enable_gqa
 
         if hasattr(SDPAParams, "enable_gqa"):
-            args.append(enable_gqa)
+            args.append(_enable_gqa)
         params = SDPAParams(query, k_and_v.keys(), k_and_v.values(), mask, 0.0, True, *args)
         if expected is SDPBackend.FLASH_ATTENTION:
             assert flash_sdp_enabled(), "flash_sdp_enabled() is False"