diff --git a/openwam/model/video_backbone/wan/shared/core/attention/attention.py b/openwam/model/video_backbone/wan/shared/core/attention/attention.py index 15803ea1..9eb491ab 100644 --- a/openwam/model/video_backbone/wan/shared/core/attention/attention.py +++ b/openwam/model/video_backbone/wan/shared/core/attention/attention.py @@ -1,8 +1,11 @@ +import logging import os import torch from einops import rearrange +logger = logging.getLogger(__name__) + try: import flash_attn_interface @@ -32,19 +35,48 @@ XFORMERS_AVAILABLE = False +def _available_implementations() -> dict: + """The implementation names this build can actually dispatch, in priority order.""" + return { + "flash_attention_3": FLASH_ATTN_3_AVAILABLE, + "flash_attention_2": FLASH_ATTN_2_AVAILABLE, + "sage_attention": SAGE_ATTN_AVAILABLE, + "xformers": XFORMERS_AVAILABLE, + "torch": True, + } + + def initialize_attention_priority(): - if os.environ.get("DIFFSYNTH_ATTENTION_IMPLEMENTATION") is not None: - return os.environ.get("DIFFSYNTH_ATTENTION_IMPLEMENTATION").lower() - elif FLASH_ATTN_3_AVAILABLE: + """Resolve the attention implementation, honouring an explicit env override. + + An override is validated the way WAM_ATTENTION_IMPL is in + openwam/model/action_backbone/components.py: an unknown name raises, and a known + name whose library did not import warns and falls back to auto-detection rather + than failing later inside the kernel wrapper. + """ + # An empty or whitespace-only value reads as "no override", matching the sibling. + override = os.environ.get("DIFFSYNTH_ATTENTION_IMPLEMENTATION", "").strip().lower() + if override: + available = _available_implementations() + if override not in available: + raise ValueError( + f"Unknown DIFFSYNTH_ATTENTION_IMPLEMENTATION='{override}'. Choose from: {sorted(available)}" + ) + if available[override]: + return override + logger.warning( + "DIFFSYNTH_ATTENTION_IMPLEMENTATION='%s' requested but not available, falling back to auto-detect", + override, + ) + if FLASH_ATTN_3_AVAILABLE: return "flash_attention_3" - elif FLASH_ATTN_2_AVAILABLE: + if FLASH_ATTN_2_AVAILABLE: return "flash_attention_2" - elif SAGE_ATTN_AVAILABLE: + if SAGE_ATTN_AVAILABLE: return "sage_attention" - elif XFORMERS_AVAILABLE: + if XFORMERS_AVAILABLE: return "xformers" - else: - return "torch" + return "torch" ATTENTION_IMPLEMENTATION = initialize_attention_priority() diff --git a/tests/test_attention_env_override.py b/tests/test_attention_env_override.py new file mode 100644 index 00000000..6ef6aa49 --- /dev/null +++ b/tests/test_attention_env_override.py @@ -0,0 +1,115 @@ +"""Regression tests for the DIFFSYNTH_ATTENTION_IMPLEMENTATION override. + +The override used to be taken verbatim, so a name whose library never imported +was returned anyway and failed later inside the kernel wrapper with a NameError. +These pin the contract against its sibling WAM_ATTENTION_IMPL in +``openwam/model/action_backbone/components.py``: unknown name raises, known but +unavailable warns and falls back, known and available is honoured. + +The availability flags are faked, so every case runs on CPU-only CI. +""" + +from __future__ import annotations + +import logging + +import pytest + +ENV = "DIFFSYNTH_ATTENTION_IMPLEMENTATION" +FLAGS = { + "flash_attention_3": "FLASH_ATTN_3_AVAILABLE", + "flash_attention_2": "FLASH_ATTN_2_AVAILABLE", + "sage_attention": "SAGE_ATTN_AVAILABLE", + "xformers": "XFORMERS_AVAILABLE", +} + + +def _shared(): + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + return shared + + +def _set_availability(monkeypatch, available: set): + """Fake which libraries imported, so these run without any of them installed.""" + shared = _shared() + for name, flag in FLAGS.items(): + monkeypatch.setattr(shared, flag, name in available) + return shared + + +def test_unknown_override_raises_and_names_the_choices(monkeypatch): + shared = _set_availability(monkeypatch, set(FLAGS)) + monkeypatch.setenv(ENV, "bogus_backend") + + with pytest.raises(ValueError) as excinfo: + shared.initialize_attention_priority() + + message = str(excinfo.value) + assert "bogus_backend" in message + # the message must be actionable, not just a rejection + for name in FLAGS: + assert name in message + + +@pytest.mark.parametrize("blank", ["", " "]) +def test_blank_override_reads_as_no_override(monkeypatch, blank): + """`DIFFSYNTH_ATTENTION_IMPLEMENTATION=` is reachable and must not be an error. + + WAM_ATTENTION_IMPL does `.strip().lower()` then `if override:`, so a blank value + falls through to auto-detection there. An accidental `export VAR=` must not turn + into a hard failure at import. + """ + shared = _set_availability(monkeypatch, {"sage_attention"}) + monkeypatch.setenv(ENV, blank) + + assert shared.initialize_attention_priority() == "sage_attention" + + +def test_unavailable_override_warns_and_falls_back(monkeypatch, caplog): + """The case that used to NameError inside the kernel wrapper.""" + shared = _set_availability(monkeypatch, {"flash_attention_2"}) + monkeypatch.setenv(ENV, "flash_attention_3") + + with caplog.at_level(logging.WARNING): + resolved = shared.initialize_attention_priority() + + assert resolved == "flash_attention_2", "should fall back to auto-detect, not honour the request" + assert "flash_attention_3" in caplog.text + assert "not available" in caplog.text + + +def test_available_override_is_honoured(monkeypatch): + shared = _set_availability(monkeypatch, {"flash_attention_2", "sage_attention"}) + monkeypatch.setenv(ENV, "sage_attention") + + assert shared.initialize_attention_priority() == "sage_attention" + + +def test_override_is_normalized(monkeypatch): + """Match the sibling, which does .strip().lower() before its membership check.""" + shared = _set_availability(monkeypatch, {"flash_attention_2"}) + monkeypatch.setenv(ENV, " FLASH_ATTENTION_2 ") + + assert shared.initialize_attention_priority() == "flash_attention_2" + + +def test_torch_is_always_selectable(monkeypatch): + """Plain SDPA needs no library, so it must be honoured even with nothing installed.""" + shared = _set_availability(monkeypatch, set()) + monkeypatch.setenv(ENV, "torch") + + assert shared.initialize_attention_priority() == "torch" + + +def test_without_override_auto_detection_is_unchanged(monkeypatch): + shared = _set_availability(monkeypatch, {"flash_attention_2", "sage_attention", "xformers"}) + monkeypatch.delenv(ENV, raising=False) + + assert shared.initialize_attention_priority() == "flash_attention_2" + + shared = _set_availability(monkeypatch, {"xformers"}) + assert shared.initialize_attention_priority() == "xformers" + + shared = _set_availability(monkeypatch, set()) + assert shared.initialize_attention_priority() == "torch"