minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-09-02 15:17:57 +08:00
committed by GitHub
parent d585cec4bd
commit 9175590aa0
10 changed files with 148 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+16 -1
View File
@@ -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
+13 -14
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/base.py
+3 -3
View File
@@ -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
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+7 -3
View File
@@ -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,
python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py
+19
View File
@@ -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
+36 -26
View File
@@ -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):
python/sglang/multimodal_gen/test/unit/test_cross_attention_backend_roles.py
+24
View File
@@ -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
+19 -5
View File
@@ -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()
python/sglang/multimodal_gen/test/unit/test_qwen3vl_text.py
+1 -1
View File
@@ -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"