[diffusion] chore: scope attention backend fallback (#34891)

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():

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
 
 

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,
         )
 

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()

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()