[diffusion] fall back to a component's default attention backend

docs/docs/sglang-diffusion/api/cli.mdx

@@ -332,7 +332,7 @@ sglang generate \
   --component-attention-backends text_encoder=torch_sdpa
 ```
 
-The component key must match a pipeline module key such as `text_encoder`, `text_encoder_2`, `transformer`, `transformer_2`, or `connectors`. Component overrides take precedence over the global `--attention-backend` only while that component is being constructed.
+The component key must match a pipeline module key such as `text_encoder`, `text_encoder_2`, `transformer`, `transformer_2`, or `connectors`. Component overrides take precedence over the global `--attention-backend` only while that component is being constructed and otherwise fail if the component cannot satisfy them. Sparse self-attention backends use a compatible dense backend for cross-attention layers. The global backend remains strict for DiT components, while auxiliary components may fall back to a compatible backend.
 
 You can also pass dotted CLI entries:
 

docs/docs/sglang-diffusion/attention_backends.mdx

@@ -9,7 +9,7 @@ This document describes the attention backends available in sglang diffusion (`s
 
 Attention backends are defined by `AttentionBackendEnum` (`sglang.multimodal_gen.runtime.platforms.interface.AttentionBackendEnum`) and selected via the CLI flag `--attention-backend`.
 
-Backend selection is performed by the shared attention layers (e.g. `LocalAttention` / `USPAttention` / `UlyssesAttention` in `sglang.multimodal_gen.runtime.layers.attention.layer`) and therefore applies to any model component using these layers (e.g. diffusion transformer / DiT and encoders).
+Backend selection is performed by the shared attention layers (e.g. `LocalAttention` / `USPAttention` / `UlyssesAttention` in `sglang.multimodal_gen.runtime.layers.attention.layer`). `--attention-backend` is strict for the diffusion transformer / DiT. Auxiliary components such as encoders and VAEs use it when compatible, then fall back to a component default or a platform-compatible backend. Use `--component-attention-backends` when an auxiliary component must use a specific backend; incompatible component overrides fail unless a sparse backend is being replaced for cross-attention.
 
 When using the diffusers backend, `--attention-backend` is passed through to diffusers'
 `set_attention_backend` (e.g., `flash`, `_flash_3_hub`, `sage`, `xformers`, `native`).
@@ -129,7 +129,14 @@ The selection order in `runtime/layers/attention/selector.py` is:
 1. `global_force_attn_backend(...)` / `global_force_attn_backend_context_manager(...)`
 2. Component override from `--component-attention-backends` while that component is being constructed
 3. CLI `--attention-backend` (`ServerArgs.attention_backend`)
-4. Auto selection (platform capability, dtype, and installed packages)
+4. Layer or component default, when declared
+5. Auto selection (platform capability, dtype, and installed packages)
+
+An explicit global backend mismatch fails for DiT self-attention. Auxiliary
+components may fall back to their declared default or another compatible backend.
+Sparse backends selected for self-attention similarly fall back to a compatible
+dense backend for cross-attention. Explicit component overrides are otherwise
+strict.
 
 ## Configuration
 
@@ -631,6 +638,10 @@ sglang generate \
 ```
 
 Component keys match pipeline module names from `model_index.json`, such as `text_encoder`, `text_encoder_2`, `transformer`, `transformer_2`, or `connectors`.
+Use this override when the fallback must be pinned: unlike the global backend,
+an incompatible component override raises an error instead of selecting another
+backend. The one role-based exception is a sparse self-attention backend, which
+uses a compatible dense backend for cross-attention layers in the same component.
 
 ### Per-request override (denoise loop)
 

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

@@ -597,6 +597,8 @@ def __init__(
         softmax_scale: float | None = None,
         causal: bool = False,
         supported_attention_backends: set[AttentionBackendEnum] | None = None,
+        default_attention_backend: AttentionBackendEnum | None = None,
+        is_cross_attention: bool = False,
         compute_dtype: torch.dtype | None = None,
         **extra_impl_args,
     ) -> None:
@@ -610,7 +612,11 @@ def __init__(
 
         dtype = compute_dtype or 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,
+            default_attention_backend=default_attention_backend,
+            is_cross_attention=is_cross_attention,
         )
         impl_cls = attn_backend.get_impl_cls()
         self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
@@ -740,7 +746,8 @@ def __init__(
               each rank's local Q shard can attend directly to the locally-held
               full KV without any collective communication.
             default_attention_backend:
-              fallback used only when no global or component override is active.
+              preferred fallback when the global backend is incompatible with
+              this layer. Explicit component overrides otherwise remain strict.
             is_cross_attention:
               sparse backend preferences may select a compatible dense backend
               for cross-attention while remaining strict for self-attention.

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

@@ -69,6 +69,7 @@ class ComponentAttnBackendContext(NamedTuple):
     backend: AttentionBackendEnum | None
     component_name: str | None
     selected_backends: dict[str, str | None]
+    allow_global_backend_fallback: bool = False
 
 
 component_attn_backend_context: ContextVar[ComponentAttnBackendContext | None] = (
@@ -113,6 +114,11 @@ def get_component_attn_backend_name() -> str | None:
     return context.component_name if context is not None else None
 
 
+def _component_allows_global_backend_fallback() -> bool:
+    context = get_component_attn_backend_context()
+    return context is not None and context.allow_global_backend_fallback
+
+
 def _record_component_attn_backend(backend_name: str, reason: str | None) -> bool:
     context = get_component_attn_backend_context()
     if context is None or context.component_name is None:
@@ -166,6 +172,7 @@ def get_attn_backend(
         )
 
     selected_backend = selected_attention_backend
+    selected_from_global_cli = False
     selection_is_explicit = selected_backend is not None
     if selected_backend is None:
         selected_backend = get_global_forced_attn_backend()
@@ -188,6 +195,7 @@ def get_attn_backend(
             selection_is_explicit = isinstance(
                 server_args, ServerArgs
             ) and server_args.is_arg_explicitly_set("attention_backend")
+            selected_from_global_cli = selection_is_explicit
 
     if selected_backend is None:
         selected_backend = default_attention_backend
@@ -197,6 +205,14 @@ def get_attn_backend(
         allowed_fallback_reason = "platform default fallback"
     elif is_cross_attention and selected_backend.is_sparse:
         allowed_fallback_reason = "dense cross-attention fallback"
+    elif selected_from_global_cli and (
+        default_attention_backend is not None
+        or _component_allows_global_backend_fallback()
+    ):
+        # The global CLI backend is strict for DiT components. Auxiliary
+        # components may instead use a declared default or platform-compatible
+        # backend. A component-specific CLI override otherwise remains strict.
+        allowed_fallback_reason = "global backend fallback"
     elif not selection_is_explicit:
         allowed_fallback_reason = "platform default fallback"
 
@@ -206,7 +222,7 @@ def get_attn_backend(
 
     candidate_backends = [selected_backend]
     if allowed_fallback_reason is not None:
-        for candidate in (None, *be_tuple):
+        for candidate in (default_attention_backend, None, *be_tuple):
             if candidate not in candidate_backends:
                 candidate_backends.append(candidate)
 
@@ -229,9 +245,17 @@ def get_attn_backend(
                 selection_error = error
             continue
 
-        candidate_name = candidate_cls.get_enum().name.lower()
+        candidate_backend = candidate_cls.get_enum()
+        candidate_name = candidate_backend.name.lower()
+        if is_cross_attention and candidate_backend.is_sparse:
+            if selection_error is None:
+                selection_error = ValueError(
+                    f"Sparse attention backend '{candidate_name}' cannot serve "
+                    "cross-attention"
+                )
+            continue
         if supported_backends and not _is_backend_supported(
-            candidate_cls.get_enum(), supported_backends
+            candidate_backend, supported_backends
         ):
             if selection_error is None:
                 selection_error = ValueError(
@@ -254,14 +278,22 @@ def get_attn_backend(
         break
 
     if attention_backend_cls is None:
+        component_name = get_component_attn_backend_name()
+        component_suffix = (
+            f" for component '{component_name}'" if component_name is not None else ""
+        )
         if unsupported_requirements:
             raise ValueError(
                 f"Attention backend '{unsupported_backend_name}' does not implement "
-                f"{', '.join(unsupported_requirements)}"
+                f"{', '.join(unsupported_requirements)}{component_suffix}"
             )
         if selection_error is not None:
-            raise selection_error
-        raise ValueError("No compatible attention backend is available")
+            raise ValueError(
+                f"{selection_error}{component_suffix}"
+            ) from selection_error
+        raise ValueError(
+            f"No compatible attention backend is available{component_suffix}"
+        )
 
     backend_name = attention_backend_cls.get_enum().name.lower()
     reason = fallback_reason
@@ -332,13 +364,19 @@ def _is_backend_supported(
 def component_attn_backend_context_manager(
     attn_backend: AttentionBackendEnum | None,
     component_name: str | None = None,
+    allow_global_backend_fallback: bool = False,
 ) -> Generator[None, None, None]:
     if attn_backend is None and component_name is None:
         yield
         return
 
     token = component_attn_backend_context.set(
-        ComponentAttnBackendContext(attn_backend, component_name, {})
+        ComponentAttnBackendContext(
+            attn_backend,
+            component_name,
+            {},
+            allow_global_backend_fallback,
+        )
     )
     try:
         yield

python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py

@@ -81,6 +81,11 @@ class ComponentLoader(ABC):
     # diffusers or transformers
     expected_library: str = ""
 
+    # --attention-backend primarily selects the DiT backend. Auxiliary
+    # components may fall back when that global choice is incompatible; an
+    # explicit --component-attention-backends entry remains strict.
+    allow_global_attention_backend_fallback = True
+
     _loaders_registered = False
 
     def __init_subclass__(cls, **kwargs):
@@ -125,9 +130,12 @@ def _load_customized_with_context(
         component_name: str,
         attn_backend: Any,
         component_attn_name: str | None,
+        allow_global_backend_fallback: bool,
     ) -> AutoModel:
         with component_attn_backend_context_manager(
-            attn_backend, component_name=component_attn_name
+            attn_backend,
+            component_name=component_attn_name,
+            allow_global_backend_fallback=allow_global_backend_fallback,
         ):
             load_kwargs = self.customized_load_kwargs_for_component(
                 server_args, component_name
@@ -144,9 +152,12 @@ def _load_native_with_context(
         transformers_or_diffusers: str,
         attn_backend: Any,
         component_attn_name: str | None,
+        allow_global_backend_fallback: bool,
     ) -> AutoModel:
         with component_attn_backend_context_manager(
-            attn_backend, component_name=component_attn_name
+            attn_backend,
+            component_name=component_attn_name,
+            allow_global_backend_fallback=allow_global_backend_fallback,
         ):
             component = self.load_native(
                 component_model_path,
@@ -198,6 +209,7 @@ def load(
                 component_name,
                 attn_backend,
                 component_attn_name,
+                self.allow_global_attention_backend_fallback,
             )
             source = "sgl-diffusion"
         except (ComponentCheckpointUnsupportedError, ComponentResidencyError):
@@ -231,6 +243,7 @@ def load(
                 transformers_or_diffusers,
                 attn_backend,
                 component_attn_name,
+                self.allow_global_attention_backend_fallback,
             )
             source = "native"
             logger.warning(
@@ -501,6 +514,10 @@ def load_customized(
 class GenericComponentLoader(ComponentLoader):
     """Generic loader for components that don't have a specific loader."""
 
+    # An unknown out-of-tree component may itself be the primary transformer.
+    # Require it to opt into fallback through a registered component loader.
+    allow_global_attention_backend_fallback = False
+
     def __init__(
         self, library="transformers", component_architecture: str | None = None
     ) -> None:
@@ -521,6 +538,8 @@ def load_component(
         transformers_or_diffusers: str,
         server_args: ServerArgs,
         component_architecture: str | None = None,
+        component_attn_backend: Any = None,
+        component_attn_name: str | None = None,
     ):
         """
         Load a pipeline component.
@@ -538,15 +557,21 @@ def load_component(
         )
 
         try:
-            # Load the component
-            return loader.load(
-                component_model_path,
-                server_args,
-                component_name,
-                transformers_or_diffusers,
-            )
-        except Exception as e:
+            with component_attn_backend_context_manager(
+                component_attn_backend,
+                component_name=component_attn_name,
+                allow_global_backend_fallback=(
+                    loader.allow_global_attention_backend_fallback
+                ),
+            ):
+                return loader.load(
+                    component_model_path,
+                    server_args,
+                    component_name,
+                    transformers_or_diffusers,
+                )
+        except Exception:
             logger.error(
                 f"Error while loading component: {component_name}, {component_model_path=}"
             )
-            raise e
+            raise

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -139,6 +139,8 @@ def _server_args_for_transformer_component(
 class TransformerLoader(ComponentLoader):
     """Shared loader for (video/audio) DiT transformers."""
 
+    allow_global_attention_backend_fallback = False
+
     component_names = [
         "transformer",
         "unconditional_transformer",

python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py

@@ -11,11 +11,9 @@
 from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import (
     LTX2ConnectorConfig,
 )
-from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
     LayerwiseOffloadableModuleMixin,
 )
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 
 
 def apply_interleaved_rotary_emb(
@@ -151,22 +149,6 @@ def __init__(
         self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
         self.to_out.append(torch.nn.Dropout(dropout))
 
-        # Scaled dot product attention
-        self.attn = USPAttention(
-            num_heads=heads,
-            head_size=self.head_dim,
-            dropout_rate=0,
-            softmax_scale=None,
-            causal=False,
-            supported_attention_backends={
-                AttentionBackendEnum.FA,
-                AttentionBackendEnum.AITER,
-                AttentionBackendEnum.TORCH_SDPA,
-                AttentionBackendEnum.SAGE_ATTN,
-                AttentionBackendEnum.SAGE_ATTN_3,
-            },
-        )
-
     def forward(
         self,
         hidden_states: torch.Tensor,

python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py

@@ -222,7 +222,7 @@ def __init__(self, dim: int, kv_dim: int, num_heads: int, eps: float = 1e-6):
             head_size=self.head_dim,
             causal=False,
             softmax_scale=None,
-            # is_cross_attention=True,
+            is_cross_attention=True,
         )
 
     def forward(

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

@@ -743,6 +743,7 @@ def __init__(
             causal=False,
             supported_attention_backends=supported_attention_backends,
             prefix=add_prefix("attn", prefix),
+            is_cross_attention=True,
         )
 
     def forward(

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

@@ -372,6 +372,7 @@ def __init__(
             head_size=self.head_dim,
             causal=False,
             skip_sequence_parallel=True,
+            is_cross_attention=True,
         )
 
     def project_kv(self, encoder_hidden_states):