From e331baaaa8167b7e6a1ed03c2e1c546f5345c92e Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 15 Aug 2026 22:00:46 +0800 Subject: [PATCH] [diffusion] chore: scope attention backend fallback (#34891) --- .../runtime/layers/attention/layer.py | 9 +- .../runtime/layers/attention/selector.py | 101 ++++++++-- .../runtime/models/dits/wanvideo.py | 1 + .../unit/test_attention_backend_selector.py | 188 ++++++++++++++++++ .../test/unit/test_wan_attention_backend.py | 36 ++++ 5 files changed, 316 insertions(+), 19 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_wan_attention_backend.py diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 352dfb853..649f233fe 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -685,6 +685,7 @@ class USPAttention(nn.Module): dropout_rate: float = 0.0, skip_sequence_parallel: bool = False, enable_packed_qkv_input_a2a: bool = False, + is_cross_attention: bool = False, **extra_impl_args, ) -> None: """ @@ -694,6 +695,9 @@ class USPAttention(nn.Module): text/image encoder outputs), the full USP pipeline is redundant: each rank's local Q shard can attend directly to the locally-held full KV without any collective communication. + is_cross_attention: + sparse backend preferences may select a compatible dense backend + for cross-attention while remaining strict for self-attention. """ super().__init__() if softmax_scale is None: @@ -706,7 +710,10 @@ class USPAttention(nn.Module): dtype = get_compute_dtype() attn_backend = get_attn_backend( - head_size, dtype, supported_attention_backends=supported_attention_backends + head_size, + dtype, + supported_attention_backends=supported_attention_backends, + is_cross_attention=is_cross_attention, ) if get_ring_parallel_world_size() > 1: if not attn_backend.supports_ring_rotation(): diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index bed284e56..634f22a95 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -17,7 +17,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionRequirements, ) from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum -from sglang.multimodal_gen.runtime.server_args import get_global_server_args +from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname @@ -153,7 +153,9 @@ def get_attn_backend( supported_attention_backends: set[AttentionBackendEnum] | None = None, selected_attention_backend: AttentionBackendEnum | None = None, attention_requirements: AttentionRequirements | None = None, + is_cross_attention: bool = False, ) -> type[AttentionBackend]: + requirements = attention_requirements or AttentionRequirements() if supported_attention_backends is None: be_tuple = tuple() else: @@ -162,9 +164,14 @@ def get_attn_backend( sorted(list(supported_attention_backends), key=lambda b: b.name) ) - selected_backend = selected_attention_backend or get_global_forced_attn_backend() + selected_backend = selected_attention_backend + selection_is_explicit = selected_backend is not None + if selected_backend is None: + selected_backend = get_global_forced_attn_backend() + selection_is_explicit = selected_backend is not None if selected_backend is None: selected_backend = get_component_forced_attn_backend() + selection_is_explicit = selected_backend is not None if selected_backend is None: server_args = get_global_server_args() if server_args.attention_backend is not None: @@ -177,30 +184,88 @@ def get_attn_backend( f"Invalid attention backend '{server_args.attention_backend}' specified via command line. " f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}" ) + selection_is_explicit = isinstance( + server_args, ServerArgs + ) and server_args.is_arg_explicitly_set("attention_backend") + + allowed_fallback_reason = None + if selected_backend is None: + allowed_fallback_reason = "platform default fallback" + elif is_cross_attention and selected_backend.is_sparse: + allowed_fallback_reason = "dense cross-attention fallback" + elif not selection_is_explicit: + allowed_fallback_reason = "platform default fallback" constraint_backend = None if selected_backend is None and len(be_tuple) == 1: constraint_backend = be_tuple[0].name.lower() - attention_backend_cls = _cached_get_attn_backend( - head_size, - dtype, - be_tuple, - selected_backend, - ) + candidate_backends = [selected_backend] + if allowed_fallback_reason is not None: + for candidate in (None, *be_tuple): + if candidate not in candidate_backends: + candidate_backends.append(candidate) + + supported_backends = set(be_tuple) + attention_backend_cls = None + fallback_reason = None + selection_error = None + unsupported_backend_name = None + unsupported_requirements = () + for candidate_index, candidate in enumerate(candidate_backends): + try: + candidate_cls = _cached_get_attn_backend( + head_size, + dtype, + be_tuple, + candidate, + ) + except ValueError as error: + if selection_error is None: + selection_error = error + continue + + candidate_name = candidate_cls.get_enum().name.lower() + if supported_backends and not _is_backend_supported( + candidate_cls.get_enum(), supported_backends + ): + if selection_error is None: + selection_error = ValueError( + f"Attention backend '{candidate_name}' is not supported by this " + f"attention layer; supported backends: " + f"{[str(backend) for backend in be_tuple]}" + ) + continue + + missing_requirements = candidate_cls.unsupported_requirements(requirements) + if missing_requirements: + if not unsupported_requirements: + unsupported_backend_name = candidate_name + unsupported_requirements = missing_requirements + continue + + attention_backend_cls = candidate_cls + if candidate_index > 0: + fallback_reason = allowed_fallback_reason + break + + if attention_backend_cls is None: + if unsupported_requirements: + raise ValueError( + f"Attention backend '{unsupported_backend_name}' does not implement " + f"{', '.join(unsupported_requirements)}" + ) + if selection_error is not None: + raise selection_error + raise ValueError("No compatible attention backend is available") backend_name = attention_backend_cls.get_enum().name.lower() - unsupported_requirements = attention_backend_cls.unsupported_requirements( - attention_requirements or AttentionRequirements() - ) - if unsupported_requirements: - raise ValueError( - f"Attention backend '{backend_name}' does not implement " - f"{', '.join(unsupported_requirements)}" - ) - reason = "component constraint" if backend_name == constraint_backend else None + reason = fallback_reason + if reason is None and backend_name == constraint_backend: + reason = "component constraint" if not _record_component_attn_backend(backend_name, reason): - logger.info_once(f"Using {backend_name} attention backend") + reason_suffix = f" ({reason})" if reason else "" + logger.info_once(f"Using {backend_name} attention backend{reason_suffix}") return attention_backend_cls diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index d50ab44a7..f62e0a736 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -204,6 +204,7 @@ class WanSelfAttention(nn.Module): causal=False, supported_attention_backends=supported_attention_backends, skip_sequence_parallel=is_cross_attention, + is_cross_attention=is_cross_attention, quant_config=quant_config, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py new file mode 100644 index 000000000..eec9bd285 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py @@ -0,0 +1,188 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( + AttentionRequirements, +) +from sglang.multimodal_gen.runtime.layers.attention.selector import ( + _cached_get_attn_backend, + get_attn_backend, +) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum +from sglang.multimodal_gen.runtime.server_args import ServerArgs + +_SELECTOR = "sglang.multimodal_gen.runtime.layers.attention.selector" + + +class _ServerArgs(ServerArgs): + def __init__(self, backend: str, *, explicit: bool) -> None: + self.attention_backend = backend + self._explicit_arg_names = {"attention_backend"} if explicit else set() + + +class _FakeSDPABackend: + @classmethod + def get_enum(cls) -> AttentionBackendEnum: + return AttentionBackendEnum.TORCH_SDPA + + @classmethod + def unsupported_requirements(cls, _requirements) -> tuple[str, ...]: + return () + + +class _FakeFABackend: + @classmethod + def get_enum(cls) -> AttentionBackendEnum: + return AttentionBackendEnum.FA + + @classmethod + def unsupported_requirements(cls, _requirements) -> tuple[str, ...]: + return () + + +class _FakeAITERBackend: + @classmethod + def get_enum(cls) -> AttentionBackendEnum: + return AttentionBackendEnum.AITER + + @classmethod + def unsupported_requirements( + cls, requirements: AttentionRequirements + ) -> tuple[str, ...]: + return ("packed varlen attention",) if requirements.packed_varlen else () + + +class _FakePlatform: + device_name = "test" + selected_backend = None + + @classmethod + def get_attn_backend_cls_str(cls, selected_backend, _head_size, _dtype): + cls.selected_backend = selected_backend + if selected_backend == AttentionBackendEnum.AITER: + return "fake.AITERBackend" + if selected_backend in (None, AttentionBackendEnum.FA): + return "fake.FABackend" + return "fake.SDPABackend" + + +_FAKE_BACKENDS = { + "fake.AITERBackend": _FakeAITERBackend, + "fake.FABackend": _FakeFABackend, + "fake.SDPABackend": _FakeSDPABackend, +} + + +class TestAttentionBackendFallback(unittest.TestCase): + def setUp(self) -> None: + _cached_get_attn_backend.cache_clear() + _FakePlatform.selected_backend = None + + def _resolve( + self, + backend: AttentionBackendEnum, + *, + explicit: bool, + is_cross_attention: bool, + supported: set[AttentionBackendEnum], + attention_requirements: AttentionRequirements | None = None, + server_args: object | None = None, + ): + if server_args is None: + server_args = _ServerArgs(backend.name.lower(), explicit=explicit) + with ( + patch(f"{_SELECTOR}.get_global_forced_attn_backend", return_value=None), + patch(f"{_SELECTOR}.get_component_forced_attn_backend", return_value=None), + patch(f"{_SELECTOR}.get_global_server_args", return_value=server_args), + patch( + "sglang.multimodal_gen.runtime.platforms.current_platform", + _FakePlatform, + ), + patch( + f"{_SELECTOR}.resolve_obj_by_qualname", + side_effect=_FAKE_BACKENDS.__getitem__, + ), + ): + return get_attn_backend( + 128, + torch.bfloat16, + supported_attention_backends=supported, + attention_requirements=attention_requirements, + is_cross_attention=is_cross_attention, + ) + + def test_implicit_platform_preference_falls_back(self): + backend = self._resolve( + AttentionBackendEnum.AITER, + explicit=False, + is_cross_attention=False, + supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, + ) + + self.assertIs(backend, _FakeFABackend) + self.assertIsNone(_FakePlatform.selected_backend) + + def test_implicit_preference_falls_back_for_missing_capability(self): + backend = self._resolve( + AttentionBackendEnum.AITER, + explicit=False, + is_cross_attention=False, + supported={AttentionBackendEnum.AITER, AttentionBackendEnum.TORCH_SDPA}, + attention_requirements=AttentionRequirements(packed_varlen=True), + ) + + self.assertIs(backend, _FakeSDPABackend) + self.assertEqual( + _FakePlatform.selected_backend, AttentionBackendEnum.TORCH_SDPA + ) + + def test_lightweight_server_args_are_treated_as_implicit(self): + backend = self._resolve( + AttentionBackendEnum.AITER, + explicit=False, + is_cross_attention=False, + supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, + server_args=SimpleNamespace(attention_backend="aiter"), + ) + + self.assertIs(backend, _FakeFABackend) + + def test_explicit_dense_mismatch_fails_closed(self): + with self.assertRaisesRegex( + ValueError, "not supported by this attention layer" + ): + self._resolve( + AttentionBackendEnum.AITER, + explicit=True, + is_cross_attention=False, + supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, + ) + + def test_sparse_backend_falls_back_for_cross_attention(self): + backend = self._resolve( + AttentionBackendEnum.LASER_ATTN, + explicit=True, + is_cross_attention=True, + supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, + ) + + self.assertIs(backend, _FakeFABackend) + self.assertIsNone(_FakePlatform.selected_backend) + + def test_sparse_backend_mismatch_fails_for_self_attention(self): + with self.assertRaisesRegex( + ValueError, "not supported by this attention layer" + ): + self._resolve( + AttentionBackendEnum.LASER_ATTN, + explicit=True, + is_cross_attention=False, + supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_wan_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_wan_attention_backend.py new file mode 100644 index 000000000..d1125813c --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_wan_attention_backend.py @@ -0,0 +1,36 @@ +import unittest +from unittest.mock import patch + +from torch import nn + +from sglang.multimodal_gen.runtime.models.dits.wanvideo import WanSelfAttention +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + +_WAN = "sglang.multimodal_gen.runtime.models.dits.wanvideo" + + +class TestWanAttentionBackendRole(unittest.TestCase): + def test_cross_attention_role_is_forwarded_to_usp(self): + with ( + patch(f"{_WAN}.ColumnParallelLinear", return_value=nn.Identity()), + patch(f"{_WAN}.RowParallelLinear", return_value=nn.Identity()), + patch(f"{_WAN}.get_tp_world_size", return_value=1), + patch(f"{_WAN}.USPAttention") as usp_attention, + ): + WanSelfAttention( + dim=128, + num_heads=1, + qk_norm=False, + is_cross_attention=True, + supported_attention_backends={ + AttentionBackendEnum.FA, + AttentionBackendEnum.TORCH_SDPA, + }, + ) + + self.assertTrue(usp_attention.call_args.kwargs["is_cross_attention"]) + self.assertTrue(usp_attention.call_args.kwargs["skip_sequence_parallel"]) + + +if __name__ == "__main__": + unittest.main()