Fix

Author: mseeger

tests/test_model.py

@@ -1470,7 +1470,8 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
                 assert can_use_flash_attention(params, True), "can_use_flash_attention(params, True) is False"
         elif expected is SDPBackend.EFFICIENT_ATTENTION:
             assert mem_efficient_sdp_enabled(), "mem_efficient_sdp_enabled() is False"
-            assert can_use_efficient_attention(params, True), "can_use_efficient_attention(params, True) is False"
+            if (not enable_gqa) or mask is None:
+                assert can_use_efficient_attention(params, True), "can_use_efficient_attention(params, True) is False"
         elif expected is SDPBackend.MATH:
             assert math_sdp_enabled(), "math_sdp_enabled() is False"
         else:
@@ -1538,7 +1539,8 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             assert can_use_flash_attention(params, True), "can_use_flash_attention(params, True) is False"
         elif expected is SDPBackend.EFFICIENT_ATTENTION:
             assert mem_efficient_sdp_enabled(), "mem_efficient_sdp_enabled() is False"
-            assert can_use_efficient_attention(params, True), "can_use_efficient_attention(params, True) is False"
+            if (not enable_gqa) or mask is None:
+                assert can_use_efficient_attention(params, True), "can_use_efficient_attention(params, True) is False"
         elif expected is SDPBackend.MATH:
             assert math_sdp_enabled(), "math_sdp_enabled() is False"
         else: