Fix

Author: mseeger

tests/test_model.py

@@ -1435,13 +1435,14 @@ def test_sdpa_choice(config):
         pytest.skip("Gemma 2 doesn't support SDPA")
 
     torch.set_default_dtype(torch.float16)
+    enable_gqa = config["n_query_groups"] < config["n_head"]
 
     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()
         if hasattr(SDPAParams, "enable_gqa"):
-            args.append(False)
+            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"
@@ -1487,13 +1488,14 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
 @torch.inference_mode()
 def test_sdpa_choice_kv_cache(config):
     torch.set_default_dtype(torch.float16)
+    enable_gqa = config["n_query_groups"] < config["n_head"]
 
     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()
         if hasattr(SDPAParams, "enable_gqa"):
-            args.append(False)
+            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()