[diffusion] refactor: admit explicit attention backends by capability (#37441)

docs/docs/sglang-diffusion/attention_backends.mdx

@@ -143,6 +143,16 @@ Sparse backends selected for self-attention similarly fall back to a compatible
 dense backend for cross-attention. Explicit component overrides are otherwise
 strict.
 
+Model and layer backend sets guide automatic selection; they are not positive
+allowlists for explicit requests. A backend selected with
+`--attention-backend` or `--component-attention-backends` can therefore be used
+without first adding it to every model's set. It still fails closed when the
+platform cannot resolve it or when it lacks a required operation such as dense
+cross-attention, packed variable-length attention, or ring-attention LSE.
+Model paths that require one backend for correctness declare it as required;
+those layers keep that backend even when the surrounding component is
+overridden.
+
 ## Configuration
 
 Some backends require additional configuration. You can pass these parameters via `--attention-backend-config`. This argument accepts:

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -289,6 +289,8 @@ def prepare_attention_backend_override(
     layer: nn.Module, target: AttentionBackendEnum
 ) -> None:
     """Build and cache the impl for ``target``; may raise, mutates nothing."""
+    if layer._required_attention_backend is not None:
+        return
     if target in layer._attn_impl_by_backend:
         return
     backend_cls = get_attn_backend(
@@ -313,6 +315,8 @@ def apply_attention_backend_override(
     layer: nn.Module, target: AttentionBackendEnum | None
 ) -> None:
     """Flip to a prepared impl (None = construction default); cannot fail."""
+    if layer._required_attention_backend is not None:
+        return
     target = target or layer._default_attn_backend
     if target is layer.backend:
         return
@@ -331,6 +335,7 @@ class UlyssesAttention(nn.Module):
         softmax_scale: float | None = None,
         causal: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        required_attention_backend: AttentionBackendEnum | None = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:
@@ -353,7 +358,10 @@ class UlyssesAttention(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,
+            selected_attention_backend=required_attention_backend,
         )
         impl_cls = attn_backend.get_impl_cls()
 
@@ -375,6 +383,7 @@ class UlyssesAttention(nn.Module):
         self._default_attn_backend = self.backend
         self._attn_impl_by_backend = {self.backend: self.attn_impl}
         self._supported_attention_backends = supported_attention_backends
+        self._required_attention_backend = required_attention_backend
         self.dtype = dtype
         self.causal = causal
         self.sp_attention_mode, self.sp_attention_mode_is_auto = (
@@ -598,6 +607,7 @@ class LocalAttention(nn.Module):
         softmax_scale: float | None = None,
         causal: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        required_attention_backend: AttentionBackendEnum | None = None,
         default_attention_backend: AttentionBackendEnum | None = None,
         is_cross_attention: bool = False,
         compute_dtype: torch.dtype | None = None,
@@ -616,6 +626,7 @@ class LocalAttention(nn.Module):
             head_size,
             dtype,
             supported_attention_backends=supported_attention_backends,
+            selected_attention_backend=required_attention_backend,
             default_attention_backend=default_attention_backend,
             is_cross_attention=is_cross_attention,
         )
@@ -638,6 +649,7 @@ class LocalAttention(nn.Module):
         self._default_attn_backend = self.backend
         self._attn_impl_by_backend = {self.backend: self.attn_impl}
         self._supported_attention_backends = supported_attention_backends
+        self._required_attention_backend = required_attention_backend
         self.dtype = dtype
 
     def forward(
@@ -731,6 +743,7 @@ class USPAttention(nn.Module):
         softmax_scale: float | None = None,
         causal: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        required_attention_backend: AttentionBackendEnum | None = None,
         default_attention_backend: AttentionBackendEnum | None = None,
         prefix: str = "",
         dropout_rate: float = 0.0,
@@ -767,6 +780,7 @@ class USPAttention(nn.Module):
             head_size,
             dtype,
             supported_attention_backends=supported_attention_backends,
+            selected_attention_backend=required_attention_backend,
             default_attention_backend=default_attention_backend,
             is_cross_attention=is_cross_attention,
         )
@@ -798,6 +812,7 @@ class USPAttention(nn.Module):
         self._default_attn_backend = self.backend
         self._attn_impl_by_backend = {self.backend: self.attn_impl}
         self._supported_attention_backends = supported_attention_backends
+        self._required_attention_backend = required_attention_backend
         self.dtype = dtype
         self.causal = causal
         self.dropout_p = dropout_rate

python/sglang/multimodal_gen/runtime/layers/attention/selector.py

@@ -221,6 +221,13 @@ def get_attn_backend(
     default_attention_backend: AttentionBackendEnum | None = None,
     is_cross_attention: bool = False,
 ) -> type[AttentionBackend]:
+    """Resolve an attention backend for one layer.
+
+    ``supported_attention_backends`` constrains automatic selection only. An
+    explicitly requested backend may be newer than a model's preference set;
+    it is admitted when the platform resolves it and the backend satisfies the
+    layer's semantic requirements.
+    """
     requirements = attention_requirements or AttentionRequirements()
     if supported_attention_backends is None:
         be_tuple = tuple()
@@ -285,7 +292,7 @@ def get_attn_backend(
             if candidate not in candidate_backends:
                 candidate_backends.append(candidate)
 
-    supported_backends = set(be_tuple)
+    automatic_backends = set(be_tuple)
     attention_backend_cls = None
     fallback_reason = None
     selection_error = None
@@ -313,8 +320,11 @@ def get_attn_backend(
                     "cross-attention"
                 )
             continue
-        if supported_backends and not _is_backend_supported(
+        explicit_candidate = selection_is_explicit and candidate_index == 0
-            candidate_backend, supported_backends
+        if (
+            automatic_backends
+            and not explicit_candidate
+            and not _is_backend_supported(candidate_backend, automatic_backends)
         ):
             if selection_error is None:
                 selection_error = ValueError(
@@ -381,17 +391,6 @@ def _cached_get_attn_backend(
         pass
     elif selected_backend is None and len(supported_attention_backends) == 1:
         selected_backend = next(iter(supported_attention_backends))
-    elif selected_backend is not None and not _is_backend_supported(
-        selected_backend, supported_attention_backends
-    ):
-        supported_attention_backends_str = [
-            supported_attention_backend.__str__()
-            for supported_attention_backend in supported_attention_backends
-        ]
-        raise ValueError(
-            f"Attention backend '{selected_backend}' is not supported by this "
-            f"attention layer; supported backends: {supported_attention_backends_str}"
-        )
 
     attention_cls = current_platform.get_attn_backend_cls_str(
         selected_backend, head_size, dtype

python/sglang/multimodal_gen/runtime/models/dits/base.py

@@ -22,9 +22,9 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
 # TODO
 class BaseDiT(nn.Module, ABC):
-    # These are runtime implementation capabilities, not checkpoint metadata.
+    # These are runtime implementation settings, not checkpoint metadata.
-    # Concrete DiT implementations override them when their tensor layout or
+    # The backend set guides automatic selection; explicit backend requests are
-    # execution semantics support only a subset of the available backends.
+    # validated against platform and layer capabilities instead of this set.
     _fsdp_shard_conditions: list = []
     _compile_conditions: list = []
     # Methods that drive a forward pass without going through __call__. FSDP2

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -735,6 +735,7 @@ class LTX2Attention(nn.Module):
         apply_gated_attention: bool = False,
         enable_packed_qkv_input_a2a: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        required_attention_backend: AttentionBackendEnum | None = None,
         prefix: str = "",
         quant_config: QuantizationConfig | None = None,
     ) -> None:
@@ -837,6 +838,7 @@ class LTX2Attention(nn.Module):
                 softmax_scale=None,
                 causal=False,
                 supported_attention_backends=supported_attention_backends,
+                required_attention_backend=required_attention_backend,
                 is_cross_attention=is_cross_attention,
                 prefix=f"{prefix}.attn",
                 enable_packed_qkv_input_a2a=self.enable_packed_qkv_input_a2a,
@@ -852,6 +854,7 @@ class LTX2Attention(nn.Module):
                 softmax_scale=None,
                 causal=False,
                 supported_attention_backends=supported_attention_backends,
+                required_attention_backend=required_attention_backend,
                 is_cross_attention=is_cross_attention,
                 prefix=f"{prefix}.attn",
                 # official LTX2 torch_sdpa uses cuDNN; cuda setup disables it
@@ -1196,10 +1199,11 @@ class LTX2TransformerBlock(nn.Module):
             use_local_attention=use_local_av_cross_attention,
             apply_gated_attention=apply_gated_attention,
             enable_packed_qkv_input_a2a=enable_packed_qkv_input_a2a,
-            supported_attention_backends=(
+            supported_attention_backends=supported_attention_backends,
-                {AttentionBackendEnum.TORCH_SDPA}
+            required_attention_backend=(
+                AttentionBackendEnum.TORCH_SDPA
                 if force_sdpa_v2a_cross_attention
-                else supported_attention_backends
+                else None
             ),
             prefix=f"{prefix}.video_to_audio_attn",
             quant_config=quant_config,

python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py

@@ -36,6 +36,7 @@ def _fake_layer(default=AttentionBackendEnum.FA) -> SimpleNamespace:
         _default_attn_backend=default,
         _attn_impl_by_backend={default: f"{default.name.lower()}_impl"},
         _supported_attention_backends=None,
+        _required_attention_backend=None,
         _attn_impl_ctor_kwargs={"num_heads": 2},
         attn_impl=f"{default.name.lower()}_impl",
         head_size=64,
@@ -223,6 +224,24 @@ class TestLayerPrepareApply(unittest.TestCase):
         self.assertEqual(layer.attn_impl, "fa_impl")
         self.assertIs(layer.backend, AttentionBackendEnum.FA)
 
+    def test_required_backend_ignores_request_override(self):
+        layer = _fake_layer(default=AttentionBackendEnum.TORCH_SDPA)
+        layer._required_attention_backend = AttentionBackendEnum.TORCH_SDPA
+
+        layer_module.prepare_attention_backend_override(
+            layer, AttentionBackendEnum.SAGE_ATTN
+        )
+        layer_module.apply_attention_backend_override(
+            layer, AttentionBackendEnum.SAGE_ATTN
+        )
+
+        self.assertEqual(
+            layer._attn_impl_by_backend,
+            {AttentionBackendEnum.TORCH_SDPA: "torch_sdpa_impl"},
+        )
+        self.assertEqual(layer.attn_impl, "torch_sdpa_impl")
+        self.assertIs(layer.backend, AttentionBackendEnum.TORCH_SDPA)
+
 
 if __name__ == "__main__":
     unittest.main()

python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py

@@ -227,15 +227,25 @@ class TestAttentionBackendFallback(unittest.TestCase):
 
         self.assertIs(backend, _FakeFABackend)
 
-    def test_explicit_dense_mismatch_fails_closed(self):
+    def test_explicit_backend_is_not_rejected_by_automatic_selection_set(self):
-        with self.assertRaisesRegex(
+        backend = self._resolve(
-            ValueError, "not supported by this attention layer"
+            AttentionBackendEnum.AITER,
-        ):
+            explicit=True,
+            is_cross_attention=False,
+            supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
+        )
+
+        self.assertIs(backend, _FakeAITERBackend)
+        self.assertEqual(_FakePlatform.selected_backend, AttentionBackendEnum.AITER)
+
+    def test_explicit_backend_still_fails_missing_capability(self):
+        with self.assertRaisesRegex(ValueError, "packed varlen attention"):
             self._resolve(
                 AttentionBackendEnum.AITER,
                 explicit=True,
                 is_cross_attention=False,
                 supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
+                attention_requirements=AttentionRequirements(packed_varlen=True),
             )
 
     def test_explicit_global_backend_uses_component_default(self):
@@ -244,6 +254,7 @@ class TestAttentionBackendFallback(unittest.TestCase):
             explicit=True,
             is_cross_attention=False,
             supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
+            attention_requirements=AttentionRequirements(packed_varlen=True),
             default_attention_backend=AttentionBackendEnum.TORCH_SDPA,
         )
 
@@ -258,24 +269,24 @@ class TestAttentionBackendFallback(unittest.TestCase):
             explicit=True,
             is_cross_attention=False,
             supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
+            attention_requirements=AttentionRequirements(packed_varlen=True),
             allow_global_backend_fallback=True,
         )
 
         self.assertIs(backend, _FakeFABackend)
         self.assertIsNone(_FakePlatform.selected_backend)
 
-    def test_explicit_component_backend_remains_strict(self):
+    def test_explicit_component_backend_ignores_automatic_selection_set(self):
-        with self.assertRaisesRegex(
+        backend = self._resolve(
-            ValueError, "not supported by this attention layer"
+            AttentionBackendEnum.FA,
-        ):
+            explicit=True,
-            self._resolve(
+            is_cross_attention=False,
-                AttentionBackendEnum.FA,
+            supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
-                explicit=True,
+            component_backend=AttentionBackendEnum.AITER,
-                is_cross_attention=False,
+            allow_global_backend_fallback=True,
-                supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
+        )
-                component_backend=AttentionBackendEnum.AITER,
+
-                allow_global_backend_fallback=True,
+        self.assertIs(backend, _FakeAITERBackend)
-            )
 
     def test_explicit_component_backend_is_consumed(self):
         backend = self._resolve(
@@ -310,16 +321,15 @@ class TestAttentionBackendFallback(unittest.TestCase):
         self.assertIs(backend, _FakeFABackend)
         self.assertIsNone(_FakePlatform.selected_backend)
 
-    def test_sparse_backend_mismatch_fails_for_self_attention(self):
+    def test_explicit_sparse_backend_is_admitted_for_self_attention(self):
-        with self.assertRaisesRegex(
+        backend = self._resolve(
-            ValueError, "not supported by this attention layer"
+            AttentionBackendEnum.LASER_ATTN,
-        ):
+            explicit=True,
-            self._resolve(
+            is_cross_attention=False,
-                AttentionBackendEnum.LASER_ATTN,
+            supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
-                explicit=True,
+        )
-                is_cross_attention=False,
+
-                supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA},
+        self.assertIs(backend, _FakeSparseBackend)
-            )
 
 
 class TestComponentAttentionBackendScope(unittest.TestCase):

python/sglang/multimodal_gen/test/unit/test_cross_attention_backend_roles.py

@@ -74,6 +74,30 @@ def test_ltx2_derives_cross_attention_role_from_context(use_local_attention):
     assert self_attention_kwargs["is_cross_attention"] is False
 
 
+@pytest.mark.parametrize("use_local_attention", [False, True])
+def test_ltx2_forwards_required_attention_backend(use_local_attention):
+    selected_layer = "LocalAttention" if use_local_attention else "USPAttention"
+    with (
+        mock.patch.object(ltx_2, "get_tp_world_size", return_value=1),
+        mock.patch.object(ltx_2, "ColumnParallelLinear", return_value=nn.Identity()),
+        mock.patch.object(ltx_2, "RowParallelLinear", return_value=nn.Identity()),
+        mock.patch.object(ltx_2, selected_layer) as attention,
+    ):
+        ltx_2.LTX2Attention(
+            query_dim=8,
+            context_dim=8,
+            heads=1,
+            dim_head=8,
+            use_local_attention=use_local_attention,
+            required_attention_backend=AttentionBackendEnum.TORCH_SDPA,
+        )
+
+    assert (
+        attention.call_args.kwargs["required_attention_backend"]
+        is AttentionBackendEnum.TORCH_SDPA
+    )
+
+
 def test_mova_bridge_marks_conditional_attention_as_cross_attention():
     with (
         mock.patch.object(mova_dual_tower, "get_tp_world_size", return_value=1),

python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py

@@ -184,17 +184,31 @@ class TestCudaAttentionBackendSelection(unittest.TestCase):
             ):
                 _SpargeAttentionBackendResolver.resolve(FakeCudaPlatform)
 
-    def test_explicit_backend_rejected_by_a_model_fails_closed(self):
+    def test_explicit_backend_is_not_rejected_by_model_preferences(self):
-        with self.assertRaisesRegex(
+        class FakeAITERBackend:
-            ValueError, "not supported by this attention layer"
+            @classmethod
+            def get_enum(cls):
+                return AttentionBackendEnum.AITER
+
+        with (
+            patch(
+                "sglang.multimodal_gen.runtime.platforms.current_platform",
+                FakeCudaPlatform,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.layers.attention.selector.resolve_obj_by_qualname",
+                return_value=FakeAITERBackend,
+            ),
         ):
-            _cached_get_attn_backend(
+            backend = _cached_get_attn_backend(
                 128,
                 torch.float16,
                 (AttentionBackendEnum.FA,),
-                AttentionBackendEnum.SAGE_ATTN,
+                AttentionBackendEnum.AITER,
             )
 
+        self.assertEqual(backend.get_enum(), AttentionBackendEnum.AITER)
+
 
 if __name__ == "__main__":
     unittest.main()

python/sglang/multimodal_gen/test/unit/test_qwen3vl_text.py

@@ -51,7 +51,7 @@ class _FakePlatform:
     @staticmethod
     def get_attn_backend_cls_str(selected_backend, _head_size, _dtype):
         if selected_backend not in (None, AttentionBackendEnum.FA):
-            raise AssertionError(f"Unexpected backend: {selected_backend}")
+            return None
         return "fake.FABackend"