diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index b5d589f79..159088d65 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -91,13 +91,13 @@ def get_attn_backend( dtype: torch.dtype, supported_attention_backends: set[AttentionBackendEnum] | None = None, ) -> type[AttentionBackend]: - if supported_attention_backends is not None: + if supported_attention_backends is None: + be_tuple = tuple() + else: # Sort the backend names to ensure consistent cache key be_tuple = tuple( sorted(list(supported_attention_backends), key=lambda b: b.name) ) - else: - be_tuple = None return _cached_get_attn_backend(head_size, dtype, be_tuple) @@ -105,7 +105,7 @@ def get_attn_backend( def _cached_get_attn_backend( head_size: int, dtype: torch.dtype, - supported_attention_backends: tuple[AttentionBackendEnum] | None = None, + supported_attention_backends: tuple[AttentionBackendEnum], ) -> type[AttentionBackend]: # Check whether a particular choice of backend was # previously forced. @@ -115,8 +115,6 @@ def _cached_get_attn_backend( from sglang.multimodal_gen.runtime.platforms import current_platform supported_attention_backends = set(supported_attention_backends) - if not supported_attention_backends: - raise ValueError("supported_attention_backends is empty") selected_backend = None backend_by_global_setting: AttentionBackendEnum | None = ( get_global_forced_attn_backend() @@ -139,12 +137,12 @@ def _cached_get_attn_backend( ) # get device-specific attn_backend - if selected_backend is None: + if len(supported_attention_backends) == 0: + # all attention backends are allowed + pass + elif selected_backend is None: logger.debug(f"Attention backend not specified") - elif ( - not supported_attention_backends - or selected_backend not in supported_attention_backends - ): + elif selected_backend not in supported_attention_backends: supported_attention_backends_str = [ supported_attention_backend.__str__() for supported_attention_backend in supported_attention_backends diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 1d54fdbbb..2ffd22e37 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux.py @@ -43,10 +43,7 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.platforms import ( - AttentionBackendEnum, - current_platform, -) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name @@ -125,13 +122,6 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin): 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( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py index ac2523164..636ff6158 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/flux_2.py @@ -29,10 +29,7 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( _apply_rotary_emb, ) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.platforms import ( - AttentionBackendEnum, - current_platform, -) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) # pylint: disable=invalid-name @@ -149,12 +146,6 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin): dropout_rate=0, softmax_scale=None, causal=False, - supported_attention_backends={ - AttentionBackendEnum.FA, - AttentionBackendEnum.TORCH_SDPA, - AttentionBackendEnum.SAGE_ATTN, - AttentionBackendEnum.SAGE_ATTN_3, - }, ) def forward( @@ -283,12 +274,6 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin): dropout_rate=0, softmax_scale=None, causal=False, - supported_attention_backends={ - AttentionBackendEnum.FA, - AttentionBackendEnum.TORCH_SDPA, - AttentionBackendEnum.SAGE_ATTN, - AttentionBackendEnum.SAGE_ATTN_3, - }, ) def forward( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 264f66492..b767b1e49 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -11,10 +11,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT -from sglang.multimodal_gen.runtime.platforms import ( - AttentionBackendEnum, - current_platform, -) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -125,10 +122,6 @@ class ZImageAttention(nn.Module): dropout_rate=0, softmax_scale=None, causal=False, - supported_attention_backends={ - AttentionBackendEnum.FA, - AttentionBackendEnum.TORCH_SDPA, - }, ) def forward( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 96a484b40..211171c42 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -64,7 +64,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( VerificationResult, ) from sglang.multimodal_gen.runtime.platforms import current_platform -from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler @@ -133,18 +132,10 @@ class DenoisingStage(PipelineStage): self.vae = vae self.pipeline = weakref.ref(pipeline) if pipeline else None + # TODO(will): hack, should use the actual one in dit self.attn_backend = get_attn_backend( head_size=attn_head_size, - dtype=torch.float16, # TODO(will): hack - supported_attention_backends={ - AttentionBackendEnum.SLIDING_TILE_ATTN, - AttentionBackendEnum.AITER, - AttentionBackendEnum.VIDEO_SPARSE_ATTN, - AttentionBackendEnum.VMOBA_ATTN, - AttentionBackendEnum.FA, - AttentionBackendEnum.TORCH_SDPA, - AttentionBackendEnum.SAGE_ATTN_3, - }, # hack + dtype=torch.float16, ) # cfg