diff --git a/docs/docs/sglang-diffusion/attention_backends.mdx b/docs/docs/sglang-diffusion/attention_backends.mdx index 23e0752b8..afa69a772 100644 --- a/docs/docs/sglang-diffusion/attention_backends.mdx +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 42506a6fd..f0776fa5d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index 49fe1ca10..8b78133fc 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/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( - candidate_backend, supported_backends + explicit_candidate = selection_is_explicit and candidate_index == 0 + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/base.py b/python/sglang/multimodal_gen/runtime/models/dits/base.py index 4d6a11af1..d58900174 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/base.py +++ b/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. - # Concrete DiT implementations override them when their tensor layout or - # execution semantics support only a subset of the available backends. + # These are runtime implementation settings, not checkpoint metadata. + # The backend set guides automatic selection; explicit backend requests are + # 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 8e25ddb05..9471d99c7 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/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=( - {AttentionBackendEnum.TORCH_SDPA} + supported_attention_backends=supported_attention_backends, + 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, diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py index 468a81d7b..83daa8316 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py +++ b/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() 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 index 9d5904a4b..b0185febe 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_backend_selector.py +++ b/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): - with self.assertRaisesRegex( - ValueError, "not supported by this attention layer" - ): + def test_explicit_backend_is_not_rejected_by_automatic_selection_set(self): + backend = self._resolve( + 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): - with self.assertRaisesRegex( - ValueError, "not supported by this attention layer" - ): - self._resolve( - AttentionBackendEnum.FA, - explicit=True, - is_cross_attention=False, - supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, - component_backend=AttentionBackendEnum.AITER, - allow_global_backend_fallback=True, - ) + def test_explicit_component_backend_ignores_automatic_selection_set(self): + backend = self._resolve( + AttentionBackendEnum.FA, + explicit=True, + is_cross_attention=False, + 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): - 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}, - ) + def test_explicit_sparse_backend_is_admitted_for_self_attention(self): + backend = self._resolve( + AttentionBackendEnum.LASER_ATTN, + explicit=True, + is_cross_attention=False, + supported={AttentionBackendEnum.FA, AttentionBackendEnum.TORCH_SDPA}, + ) + + self.assertIs(backend, _FakeSparseBackend) class TestComponentAttentionBackendScope(unittest.TestCase): diff --git a/python/sglang/multimodal_gen/test/unit/test_cross_attention_backend_roles.py b/python/sglang/multimodal_gen/test/unit/test_cross_attention_backend_roles.py index 84d4e0243..e120a6a72 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cross_attention_backend_roles.py +++ b/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), diff --git a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py index 603e95cb5..fbf7522b7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py +++ b/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): - with self.assertRaisesRegex( - ValueError, "not supported by this attention layer" + def test_explicit_backend_is_not_rejected_by_model_preferences(self): + class FakeAITERBackend: + @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() diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen3vl_text.py b/python/sglang/multimodal_gen/test/unit/test_qwen3vl_text.py index d7cf6fb6d..908e99d4e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen3vl_text.py +++ b/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"