From 5ddde5cf7b878efd23f93bda322081a4e1708aca Mon Sep 17 00:00:00 2001 From: wayrise Date: Wed, 9 Sep 2026 00:26:44 +0800 Subject: [PATCH] fix(wan/shared): validate DIFFSYNTH_ATTENTION_IMPLEMENTATION initialize_attention_priority() lowercased the override and returned it unchecked, unlike its sibling WAM_ATTENTION_IMPL, which validates against _BACKEND_MAP and warns when the requested library is missing. Three consequences, all reported in #12: - Naming an uninstalled backend crashed at dispatch. flash_attention_3() dereferences the module-level flash_attn_interface that the try/except import never bound, so DIFFSYNTH_ATTENTION_IMPLEMENTATION=flash_attention_3 on a box without FA3 raised NameError from inside the kernel wrapper. - A typo was accepted silently. attention_forward's else-branch ran it as plain SDPA, so a misspelled backend was indistinguishable from a working one and a whole job could train on SDPA while looking like it had a fused kernel. - The empty string was reachable, because the lookup guarded on `is not None` rather than truthiness, and behaved the same way. It is also why the deploy diagnostics logged a blank backend name. Mirror the sibling: strip and lowercase, raise on a name outside KNOWN_ATTENTION_IMPLEMENTATIONS, and fall back to auto-detection with a warning when the named backend's library did not import. Auto-detection order is unchanged and pinned by a test. Rejecting an unknown name raises at import, since ATTENTION_IMPLEMENTATION is resolved at module scope. That is deliberate: a typo should stop the run rather than quietly downgrade it, and the message names the valid choices. _log_attention_backends already catches the import and reports it per subsystem. Closes #12. Co-Authored-By: Claude Opus 5 --- .../wan/shared/core/attention/attention.py | 63 ++++++++++--- tests/test_attention_backend_dispatch.py | 93 +++++++++++++++++++ 2 files changed, 144 insertions(+), 12 deletions(-) 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..62557cbb 100644 --- a/openwam/model/video_backbone/wan/shared/core/attention/attention.py +++ b/openwam/model/video_backbone/wan/shared/core/attention/attention.py @@ -1,3 +1,4 @@ +import logging import os import torch @@ -32,19 +33,57 @@ XFORMERS_AVAILABLE = False +logger = logging.getLogger(__name__) + +# Highest priority first. "torch" is always available and terminates auto-detection. +KNOWN_ATTENTION_IMPLEMENTATIONS = ( + "flash_attention_3", + "flash_attention_2", + "sage_attention", + "xformers", + "torch", +) + + +def _implementation_available(implementation: str) -> bool: + """Whether the library backing ``implementation`` actually imported.""" + 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, + }[implementation] + + 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: - return "flash_attention_3" - elif FLASH_ATTN_2_AVAILABLE: - return "flash_attention_2" - elif SAGE_ATTN_AVAILABLE: - return "sage_attention" - elif XFORMERS_AVAILABLE: - return "xformers" - else: - return "torch" + """Pick the attention implementation, honouring DIFFSYNTH_ATTENTION_IMPLEMENTATION. + + The 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 is missing warns and falls back to auto-detection. Returning it + unchecked let a typo run a whole job on SDPA while looking like a fused backend, and + let a name like "flash_attention_3" reach the kernel wrapper and raise NameError. + """ + override = os.environ.get("DIFFSYNTH_ATTENTION_IMPLEMENTATION", "").strip().lower() + + if override: + if override not in KNOWN_ATTENTION_IMPLEMENTATIONS: + raise ValueError( + f"Unknown DIFFSYNTH_ATTENTION_IMPLEMENTATION='{override}'. " + f"Choose from: {list(KNOWN_ATTENTION_IMPLEMENTATIONS)}" + ) + if _implementation_available(override): + return override + logger.warning( + "DIFFSYNTH_ATTENTION_IMPLEMENTATION='%s' requested but not available, falling back to auto-detect", + override, + ) + + for implementation in KNOWN_ATTENTION_IMPLEMENTATIONS: + if _implementation_available(implementation): + return implementation + return "torch" ATTENTION_IMPLEMENTATION = initialize_attention_priority() diff --git a/tests/test_attention_backend_dispatch.py b/tests/test_attention_backend_dispatch.py index 5366b810..833c0b19 100644 --- a/tests/test_attention_backend_dispatch.py +++ b/tests/test_attention_backend_dispatch.py @@ -319,3 +319,96 @@ def test_wan_backend_names_match_their_sources(): assert set(server._WAN_BACKENDS) == fused | {"xformers"}, ( "a Wan backend was added or removed; classify it in _WAN_BACKENDS" ) + + +# ------------------------------------------- DIFFSYNTH_ATTENTION_IMPLEMENTATION validation + + +_AVAILABILITY_FLAGS = ( + "FLASH_ATTN_3_AVAILABLE", + "FLASH_ATTN_2_AVAILABLE", + "SAGE_ATTN_AVAILABLE", + "XFORMERS_AVAILABLE", +) + + +def _reinitialise(monkeypatch, shared, override, **available): + """Re-run initialize_attention_priority() under a given env value and flag set.""" + if override is None: + monkeypatch.delenv("DIFFSYNTH_ATTENTION_IMPLEMENTATION", raising=False) + else: + monkeypatch.setenv("DIFFSYNTH_ATTENTION_IMPLEMENTATION", override) + for flag in _AVAILABILITY_FLAGS: + monkeypatch.setattr(shared, flag, available.get(flag, False)) + return shared.initialize_attention_priority() + + +def test_diffsynth_override_rejects_an_unknown_value(monkeypatch): + """A typo must not run a whole job on SDPA while reading like a fused backend.""" + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + with pytest.raises(ValueError, match="bogus"): + _reinitialise(monkeypatch, shared, "bogus") + + +def test_diffsynth_override_falls_back_when_the_library_is_missing(monkeypatch, caplog): + """Naming an uninstalled backend used to reach the kernel wrapper and raise NameError. + + flash_attention_3() dereferences the module-level flash_attn_interface, which the + try/except import never bound, so dispatch died with a NameError rather than a + diagnosable message. The sibling WAM_ATTENTION_IMPL has always warned and fallen back. + """ + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + with caplog.at_level(logging.WARNING): + implementation = _reinitialise(monkeypatch, shared, "flash_attention_3", FLASH_ATTN_2_AVAILABLE=True) + + assert implementation == "flash_attention_2" + assert "flash_attention_3" in caplog.text + + +def test_diffsynth_override_never_returns_an_unavailable_backend(monkeypatch): + """With nothing installed, every override degrades to plain torch.""" + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + for override in shared.KNOWN_ATTENTION_IMPLEMENTATIONS: + assert _reinitialise(monkeypatch, shared, override) == "torch" + + +def test_diffsynth_override_is_honoured_when_available(monkeypatch): + """A valid, installed override still wins over the auto-detection priority.""" + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + picked = _reinitialise( + monkeypatch, + shared, + "sage_attention", + SAGE_ATTN_AVAILABLE=True, + FLASH_ATTN_2_AVAILABLE=True, # higher priority, but not what was asked for + ) + assert picked == "sage_attention" + + +def test_diffsynth_override_ignores_blank_values_and_normalises_the_name(monkeypatch): + """An empty override is reachable: the old lookup guarded on `is not None`, not truthiness.""" + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + for blank in ("", " "): + assert _reinitialise(monkeypatch, shared, blank, FLASH_ATTN_2_AVAILABLE=True) == "flash_attention_2" + + assert _reinitialise(monkeypatch, shared, " SAGE_ATTENTION ", SAGE_ATTN_AVAILABLE=True) == "sage_attention" + + +def test_diffsynth_auto_detection_priority_is_unchanged(monkeypatch): + """Pin the pre-existing order, so validation did not quietly reshuffle it.""" + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + expected = [ + ({"FLASH_ATTN_3_AVAILABLE", "FLASH_ATTN_2_AVAILABLE", "SAGE_ATTN_AVAILABLE"}, "flash_attention_3"), + ({"FLASH_ATTN_2_AVAILABLE", "SAGE_ATTN_AVAILABLE", "XFORMERS_AVAILABLE"}, "flash_attention_2"), + ({"SAGE_ATTN_AVAILABLE", "XFORMERS_AVAILABLE"}, "sage_attention"), + ({"XFORMERS_AVAILABLE"}, "xformers"), + (set(), "torch"), + ] + for installed, winner in expected: + assert _reinitialise(monkeypatch, shared, None, **{f: True for f in installed}) == winner