migrate remaining consumers to canonical attn_implementation

2026-04-23 21:26:18 +00:00
parent bce65e3332
commit a0d24bcc19
3 changed files with 81 additions and 3 deletions
--- a/tests/e2e/multigpu/test_llama.py
+++ b/tests/e2e/multigpu/test_llama.py
@@ -521,9 +521,9 @@ class TestMultiGPULlama:
            }
        )
        if attention_backend == "flash":
-            cfg.flash_attention = True
+            cfg.attn_implementation = "flash_attention_2"
        elif attention_backend == "flex":
-            cfg.flex_attention = True
+            cfg.attn_implementation = "flex_attention"

        # write cfg to yaml file
        Path(temp_dir).mkdir(parents=True, exist_ok=True)