Fix

Author: mseeger

tests/test_model.py

@@ -1455,14 +1455,15 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             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)
+            _k_and_v = DefaultKeysAndValues(key, value)
             _enable_gqa = False
         else:
             _enable_gqa = enable_gqa
+            _k_and_v = k_and_v
 
         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)
+        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"
             if config.sliding_window_size is None:
@@ -1523,14 +1524,15 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
             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)
+            _k_and_v = DefaultKeysAndValues(key, value)
             _enable_gqa = False
         else:
             _enable_gqa = enable_gqa
+            _k_and_v = k_and_v
 
         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)
+        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"
             assert can_use_flash_attention(params, True), "can_use_flash_attention(params, True) is False"