From b9f0b1fa31ce8e9c01f5528d376bd0e910fa693a Mon Sep 17 00:00:00 2001 From: wayrise Date: Tue, 8 Sep 2026 17:58:22 +0800 Subject: [PATCH] fix(attention): guard fused backend dispatch on device and dtype #7 fixed flash_attention() in the Wan video DiT. The same defect is present at the other dispatch sites: FA2/FA3/sage are CUDA half-precision kernels that raise rather than degrade, so selecting one from import availability alone breaks every fp32 or CPU call once flash-attn is installed. - action_backbone/components.py: the guards checked device but not dtype, so cuda/fp32 still raised here. xformers stays device-only, since memory_efficient_attention does support fp32 and guarding it on half precision would downgrade it needlessly. - video_backbone/wan/shared/core/attention/attention.py: resolve the configured implementation per call instead of once at import time. - deploy/server.py: reported the DiT backend from the availability flags alone, so it could log flash_attention_2 for a run that correctly used SDPA. It now asks dit.py, through a new fused_backend_name() helper, and qualifies the answer. The tests fake the backends, so they cover the "flash-attn is installed" configuration on CPU-only CI, which is where this class of bug hides. They cover the dit.py path from #7 too, which merged without a regression test. Follow-up to #7, reported by @rakhimovv. Co-Authored-By: Claude Opus 5 --- openwam/deploy/server.py | 18 +- openwam/model/action_backbone/components.py | 11 +- .../model/video_backbone/wan/models/dit.py | 11 + .../wan/shared/core/attention/attention.py | 24 +- tests/test_attention_backend_dispatch.py | 214 ++++++++++++++++++ 5 files changed, 262 insertions(+), 16 deletions(-) create mode 100644 tests/test_attention_backend_dispatch.py diff --git a/openwam/deploy/server.py b/openwam/deploy/server.py index e8fa31be..3ee04b74 100644 --- a/openwam/deploy/server.py +++ b/openwam/deploy/server.py @@ -376,14 +376,9 @@ def _log_attention_backends(logger): try: import openwam.model.video_backbone.wan.models.dit as _vdit - if getattr(_vdit, "FLASH_ATTN_3_AVAILABLE", False): - vdit_backend = "flash_attention_3" - elif getattr(_vdit, "FLASH_ATTN_2_AVAILABLE", False): - vdit_backend = "flash_attention_2" - elif getattr(_vdit, "SAGE_ATTN_AVAILABLE", False): - vdit_backend = "sage_attention" - else: - vdit_backend = "torch_sdpa" + vdit_backend = _vdit.fused_backend_name() + if vdit_backend != "torch_sdpa": + vdit_backend += " (CUDA fp16/bf16 only; torch_sdpa otherwise)" lines.append(f" Video DiT : {vdit_backend}") except Exception as e: lines.append(f" Video DiT : ERROR ({e})") @@ -392,7 +387,12 @@ def _log_attention_backends(logger): try: from openwam.model.video_backbone.wan.shared.core.attention.attention import ATTENTION_IMPLEMENTATION - lines.append(f" Wan shared core : {ATTENTION_IMPLEMENTATION}") + shared_backend = ATTENTION_IMPLEMENTATION + if shared_backend == "xformers": + shared_backend += " (CUDA only; torch_sdpa otherwise)" + elif shared_backend != "torch": + shared_backend += " (CUDA fp16/bf16 only; torch_sdpa otherwise)" + lines.append(f" Wan shared core : {shared_backend}") except Exception as e: lines.append(f" Wan shared core : ERROR ({e})") diff --git a/openwam/model/action_backbone/components.py b/openwam/model/action_backbone/components.py index 712d9423..c59a5785 100644 --- a/openwam/model/action_backbone/components.py +++ b/openwam/model/action_backbone/components.py @@ -272,12 +272,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: _ATTENTION_FN: Callable | None = None +def _fused_kernel_usable(q: Tensor) -> bool: + """FA2/FA3/sage are CUDA half-precision kernels; they raise on anything else.""" + return q.is_cuda and q.dtype in (torch.float16, torch.bfloat16) + + def _try_flash_attn_3() -> Callable | None: try: from flash_attn_interface import flash_attn_func as flash3_fn # type: ignore def _flash3(q: Tensor, k: Tensor, v: Tensor) -> Tensor: - if q.device.type != "cuda": + if not _fused_kernel_usable(q): return _sdpa(q, k, v) q = q.transpose(1, 2) k = k.transpose(1, 2) @@ -295,7 +300,7 @@ def _try_flash_attn_2() -> Callable | None: from flash_attn import flash_attn_func # type: ignore def _flash2(q: Tensor, k: Tensor, v: Tensor) -> Tensor: - if q.device.type != "cuda": + if not _fused_kernel_usable(q): return _sdpa(q, k, v) q = q.transpose(1, 2) k = k.transpose(1, 2) @@ -313,7 +318,7 @@ def _try_sage_attention() -> Callable | None: from sageattention import sageattn # type: ignore def _sage(q: Tensor, k: Tensor, v: Tensor) -> Tensor: - if q.device.type != "cuda": + if not _fused_kernel_usable(q): return _sdpa(q, k, v) return sageattn(q, k, v) diff --git a/openwam/model/video_backbone/wan/models/dit.py b/openwam/model/video_backbone/wan/models/dit.py index 4ea6805c..37aabb1d 100644 --- a/openwam/model/video_backbone/wan/models/dit.py +++ b/openwam/model/video_backbone/wan/models/dit.py @@ -31,6 +31,17 @@ SAGE_ATTN_AVAILABLE = False +def fused_backend_name() -> str: + """Fused backend used for CUDA fp16/bf16 input. Anything else falls back to SDPA.""" + if FLASH_ATTN_3_AVAILABLE: + return "flash_attention_3" + if FLASH_ATTN_2_AVAILABLE: + return "flash_attention_2" + if SAGE_ATTN_AVAILABLE: + return "sage_attention" + return "torch_sdpa" + + def flash_attention( q: torch.Tensor, k: torch.Tensor, 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 670eaccf..15803ea1 100644 --- a/openwam/model/video_backbone/wan/shared/core/attention/attention.py +++ b/openwam/model/video_backbone/wan/shared/core/attention/attention.py @@ -170,6 +170,21 @@ def xformers_attention( return out +def resolve_implementation(q: torch.Tensor) -> str: + """Narrow ATTENTION_IMPLEMENTATION to one this input can actually run. + + FA2/FA3/sage are CUDA half-precision kernels and xformers is CUDA-only; all of + them raise rather than degrade, so anything they cannot take falls back to SDPA. + xformers keeps its fp32 support, hence the separate check. + """ + if ATTENTION_IMPLEMENTATION in ("flash_attention_3", "flash_attention_2", "sage_attention"): + if not (q.is_cuda and q.dtype in (torch.float16, torch.bfloat16)): + return "torch" + elif ATTENTION_IMPLEMENTATION == "xformers" and not q.is_cuda: + return "torch" + return ATTENTION_IMPLEMENTATION + + def attention_forward( q: torch.Tensor, k: torch.Tensor, @@ -186,13 +201,14 @@ def attention_forward( if compatibility_mode or (attn_mask is not None): return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, attn_mask=attn_mask, scale=scale) else: - if ATTENTION_IMPLEMENTATION == "flash_attention_3": + impl = resolve_implementation(q) + if impl == "flash_attention_3": return flash_attention_3(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) - elif ATTENTION_IMPLEMENTATION == "flash_attention_2": + elif impl == "flash_attention_2": return flash_attention_2(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) - elif ATTENTION_IMPLEMENTATION == "sage_attention": + elif impl == "sage_attention": return sage_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) - elif ATTENTION_IMPLEMENTATION == "xformers": + elif impl == "xformers": return xformers_attention(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) else: return torch_sdpa(q, k, v, q_pattern, k_pattern, v_pattern, out_pattern, dims, scale=scale) diff --git a/tests/test_attention_backend_dispatch.py b/tests/test_attention_backend_dispatch.py new file mode 100644 index 00000000..640bf563 --- /dev/null +++ b/tests/test_attention_backend_dispatch.py @@ -0,0 +1,214 @@ +"""Regression tests for fused-attention backend dispatch. + +FA2/FA3/sage are CUDA half-precision kernels: they raise rather than degrade. +Every dispatch site must therefore check device *and* dtype before selecting +one, otherwise merely having ``flash-attn`` importable breaks fp32 and CPU +inputs -- including the CPU tensors the rest of this suite runs on. + +The backends are faked here, so these tests exercise the "flash-attn is +installed" configuration without it being installed. That matters: CI runs on +CPU-only torch and would otherwise never cover this dispatch at all. +""" + +from __future__ import annotations + +import logging +import types + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange + + +class _FakeQuery: + """Stand-in exposing only the two attributes the dispatch predicates read. + + Lets the CUDA-side branches be covered on a CPU-only box. + """ + + def __init__(self, is_cuda: bool, dtype: torch.dtype): + self.is_cuda = is_cuda + self.dtype = dtype + + +def _sdpa_reference(q, k, v, num_heads): + qq, kk, vv = (rearrange(t, "b s (n d) -> b n s d", n=num_heads) for t in (q, k, v)) + return rearrange(F.scaled_dot_product_attention(qq, kk, vv), "b n s d -> b s (n d)", n=num_heads) + + +def _install_fake_flash_attn_2(monkeypatch, vdit, kernel): + """Make ``dit.flash_attention`` believe FA2 -- and only FA2 -- is installed.""" + monkeypatch.setattr(vdit, "flash_attn", types.SimpleNamespace(flash_attn_func=kernel), raising=False) + monkeypatch.setattr(vdit, "FLASH_ATTN_2_AVAILABLE", True) + monkeypatch.setattr(vdit, "FLASH_ATTN_3_AVAILABLE", False) + monkeypatch.setattr(vdit, "SAGE_ATTN_AVAILABLE", False) + + +# --------------------------------------------------------------- Wan video DiT + + +def test_wan_dit_falls_back_to_sdpa_on_cpu(monkeypatch): + """CPU tensors must never reach the fused kernel.""" + import openwam.model.video_backbone.wan.models.dit as vdit + + def _reject(*args, **kwargs): + raise AssertionError("fused kernel must not receive CPU tensors") + + _install_fake_flash_attn_2(monkeypatch, vdit, _reject) + + num_heads = 4 + q, k, v = (torch.randn(2, 8, num_heads * 16) for _ in range(3)) + + out = vdit.flash_attention(q, k, v, num_heads=num_heads) + + assert torch.allclose(out, _sdpa_reference(q, k, v, num_heads), atol=0, rtol=0) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_wan_dit_falls_back_to_sdpa_for_cuda_fp32(monkeypatch): + """fp32 is rejected by the fused kernels even on CUDA.""" + import openwam.model.video_backbone.wan.models.dit as vdit + + def _reject(*args, **kwargs): + raise AssertionError("fused kernel must not receive fp32 tensors") + + _install_fake_flash_attn_2(monkeypatch, vdit, _reject) + + num_heads = 4 + q, k, v = (torch.randn(2, 8, num_heads * 16, device="cuda", dtype=torch.float32) for _ in range(3)) + + out = vdit.flash_attention(q, k, v, num_heads=num_heads) + + assert torch.allclose(out, _sdpa_reference(q, k, v, num_heads), atol=0, rtol=0) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_wan_dit_keeps_fused_path_for_cuda_half(monkeypatch): + """The guard must not cost us the fast path it is protecting.""" + import openwam.model.video_backbone.wan.models.dit as vdit + + calls = [] + + def _record(q, k, v, *args, **kwargs): + calls.append(1) + return q + + _install_fake_flash_attn_2(monkeypatch, vdit, _record) + + num_heads = 4 + q, k, v = (torch.randn(2, 8, num_heads * 16, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + + out = vdit.flash_attention(q, k, v, num_heads=num_heads) + + assert calls, "cuda/bf16 input should still reach the fused kernel" + assert torch.equal(out, q) + + +# ------------------------------------------------------------------ ActionDiT + + +def test_action_fused_kernel_usable_requires_cuda_and_half(): + from openwam.model.action_backbone.components import _fused_kernel_usable + + assert _fused_kernel_usable(_FakeQuery(True, torch.float16)) + assert _fused_kernel_usable(_FakeQuery(True, torch.bfloat16)) + # the case the device-only guard used to let through + assert not _fused_kernel_usable(_FakeQuery(True, torch.float32)) + assert not _fused_kernel_usable(_FakeQuery(False, torch.float32)) + assert not _fused_kernel_usable(_FakeQuery(False, torch.bfloat16)) + + +def test_action_sage_backend_falls_back_on_cpu(monkeypatch): + """Same contract as the flash-attn backends, for the sage path.""" + import sys + + from openwam.model.action_backbone import components + + def _cuda_only_sageattn(*args, **kwargs): + raise AssertionError("sageattn should not receive CPU tensors") + + monkeypatch.setitem(sys.modules, "sageattention", types.SimpleNamespace(sageattn=_cuda_only_sageattn)) + + fn = components._try_sage_attention() + assert fn is not None + + q, k, v = (torch.randn(1, 2, 3, 4) for _ in range(3)) + assert fn(q, k, v).shape == q.shape + + +# ----------------------------------------------------------- Wan shared core + + +@pytest.mark.parametrize("impl", ["flash_attention_3", "flash_attention_2", "sage_attention"]) +def test_shared_core_downgrades_fused_backends_it_cannot_run(monkeypatch, impl): + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + monkeypatch.setattr(shared, "ATTENTION_IMPLEMENTATION", impl) + + assert shared.resolve_implementation(_FakeQuery(True, torch.bfloat16)) == impl + assert shared.resolve_implementation(_FakeQuery(True, torch.float32)) == "torch" + assert shared.resolve_implementation(_FakeQuery(False, torch.float32)) == "torch" + + +def test_shared_core_keeps_xformers_for_cuda_fp32(monkeypatch): + """xformers is CUDA-only but does support fp32, so it must not be downgraded for it.""" + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + monkeypatch.setattr(shared, "ATTENTION_IMPLEMENTATION", "xformers") + + assert shared.resolve_implementation(_FakeQuery(True, torch.float32)) == "xformers" + assert shared.resolve_implementation(_FakeQuery(False, torch.float32)) == "torch" + + +def test_shared_core_attention_forward_runs_on_cpu(monkeypatch): + """End-to-end: a CPU call must produce the SDPA result, not raise.""" + from openwam.model.video_backbone.wan.shared.core.attention import attention as shared + + monkeypatch.setattr(shared, "ATTENTION_IMPLEMENTATION", "flash_attention_2") + monkeypatch.setattr(shared, "flash_attn_interface", None, raising=False) + + q, k, v = (torch.randn(2, 4, 8, 16) for _ in range(3)) + + out = shared.attention_forward(q, k, v) + + assert torch.allclose(out, F.scaled_dot_product_attention(q, k, v), atol=0, rtol=0) + + +# ------------------------------------------------------- deploy diagnostics + + +def test_fused_backend_name_follows_availability(monkeypatch): + import openwam.model.video_backbone.wan.models.dit as vdit + + for flag in ("FLASH_ATTN_3_AVAILABLE", "FLASH_ATTN_2_AVAILABLE", "SAGE_ATTN_AVAILABLE"): + monkeypatch.setattr(vdit, flag, False) + assert vdit.fused_backend_name() == "torch_sdpa" + + monkeypatch.setattr(vdit, "SAGE_ATTN_AVAILABLE", True) + assert vdit.fused_backend_name() == "sage_attention" + + monkeypatch.setattr(vdit, "FLASH_ATTN_2_AVAILABLE", True) + assert vdit.fused_backend_name() == "flash_attention_2" + + monkeypatch.setattr(vdit, "FLASH_ATTN_3_AVAILABLE", True) + assert vdit.fused_backend_name() == "flash_attention_3" + + +def test_diagnostics_do_not_claim_a_backend_unconditionally(monkeypatch, caplog): + """The report must not read as 'flash_attention_2' on a run that uses SDPA.""" + import openwam.model.video_backbone.wan.models.dit as vdit + from openwam.deploy.server import _log_attention_backends + + monkeypatch.setattr(vdit, "FLASH_ATTN_2_AVAILABLE", True) + monkeypatch.setattr(vdit, "FLASH_ATTN_3_AVAILABLE", False) + monkeypatch.setattr(vdit, "SAGE_ATTN_AVAILABLE", False) + + with caplog.at_level(logging.INFO): + _log_attention_backends(logging.getLogger(__name__)) + + video_line = next(line for line in caplog.text.splitlines() if "Video DiT" in line) + assert "flash_attention_2" in video_line + assert "CUDA fp16/bf16 only" in video_line