Fix

Author: mseeger

tests/test_model.py

@@ -1520,10 +1520,12 @@ def assert_sdpa_backend(original_fn, query, k_and_v, mask, return_scores):
         # best effort, if the GPU can load it
         pytest.xfail()
 
-    model.mha.scaled_dot_product_attention = partial(
-        assert_sdpa_backend,
-        model.mha.scaled_dot_product_attention,
-    )
+    for block in model.transformer.h:
+        kv_cache = block.attn.kv_cache
+        kv_cache.mha.scaled_dot_product_attention = partial(
+            assert_sdpa_backend,
+            kv_cache.mha.scaled_dot_product_attention,
+        )
 
     if SUPPORTS_FLASH_ATTENTION:
         # flash attention does not support an attention mask