diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py index a6d341c64..efd69a1c6 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py @@ -11,6 +11,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionMetadata, AttentionMetadataBuilder, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum class AITerBackend(AttentionBackend): @@ -19,8 +20,8 @@ class AITerBackend(AttentionBackend): """ @staticmethod - def get_name() -> str: - return "AITER" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.AITER @staticmethod def get_impl_cls() -> type["AITerImpl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 3463ef05c..ab1293d2a 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py @@ -12,6 +12,8 @@ if TYPE_CHECKING: import torch +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + class AttentionBackend(ABC): """Abstract class for attention backends.""" @@ -23,7 +25,7 @@ class AttentionBackend(ABC): @staticmethod @abstractmethod - def get_name() -> str: + def get_enum() -> AttentionBackendEnum: raise NotImplementedError @staticmethod diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index 2394de5fb..adbe410a9 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -7,6 +7,7 @@ from typing import Any import torch from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum try: from sgl_kernel.flash_attn import flash_attn_varlen_func @@ -73,8 +74,8 @@ class FlashAttentionBackend(AttentionBackend): return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod - def get_name() -> str: - return "FLASH_ATTN" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.FA @staticmethod def get_impl_cls() -> type["FlashAttentionImpl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py index df795e062..62a1974ad 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py @@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( flash_attn_func, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -26,8 +27,8 @@ class FlashAttention2Backend(AttentionBackend): return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod - def get_name() -> str: - return "FA" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.FA2 @staticmethod def get_impl_cls() -> type["FlashAttention2Impl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py index 84d3814c4..db4643096 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py @@ -11,6 +11,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionImpl, AttentionMetadata, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -24,8 +25,8 @@ class SageAttentionBackend(AttentionBackend): return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod - def get_name() -> str: - return "SAGE_ATTN" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.SAGE_ATTN @staticmethod def get_impl_cls() -> type["SageAttentionImpl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py index c49d21525..e23588f88 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py @@ -10,6 +10,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionImpl, AttentionMetadata, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -24,8 +25,8 @@ class SageAttention3Backend(AttentionBackend): return [64, 128, 256] @staticmethod - def get_name() -> str: - return "SAGE_ATTN_3" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.SAGE_ATTN_3 @staticmethod def get_impl_cls() -> type["SageAttention3Impl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index bfa3b430d..5523f6559 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -9,6 +9,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionImpl, AttentionMetadata, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -23,8 +24,8 @@ class SDPABackend(AttentionBackend): return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod - def get_name() -> str: - return "SDPA" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.TORCH_SDPA @staticmethod def get_impl_cls() -> type["SDPAImpl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py index 6db3785ff..e89392ae3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py @@ -20,6 +20,7 @@ from sglang.multimodal_gen.runtime.managers.forward_context import ( ForwardContext, get_forward_context, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import dict_to_3d_list @@ -55,8 +56,8 @@ class SlidingTileAttentionBackend(AttentionBackend): return [32, 64, 96, 128, 160, 192, 224, 256] @staticmethod - def get_name() -> str: - return "SLIDING_TILE_ATTN" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.SLIDING_TILE_ATTN @staticmethod def get_impl_cls() -> type["SlidingTileAttentionImpl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py index 6fe342922..2ee9a17df 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py @@ -21,6 +21,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionMetadata, AttentionMetadataBuilder, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -131,8 +132,8 @@ class VideoSparseAttentionBackend(AttentionBackend): return [64, 128] @staticmethod - def get_name() -> str: - return "VIDEO_SPARSE_ATTN" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.VIDEO_SPARSE_ATTN @staticmethod def get_impl_cls() -> type["VideoSparseAttentionImpl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py index 5709601d2..e07c74336 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py @@ -19,6 +19,7 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i AttentionMetadata, AttentionMetadataBuilder, ) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -29,8 +30,8 @@ class VMOBAAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @staticmethod - def get_name() -> str: - return "VMOBA_ATTN" + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.VMOBA_ATTN @staticmethod def get_impl_cls() -> type["VMOBAAttentionImpl"]: diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index df4f377df..26bcbdd48 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -20,10 +20,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( AttentionImpl, ) -from sglang.multimodal_gen.runtime.layers.attention.selector import ( - backend_name_to_enum, - get_attn_backend, -) +from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend from sglang.multimodal_gen.runtime.layers.usp import ( _usp_input_all_to_all, _usp_output_all_to_all, @@ -78,7 +75,7 @@ class UlyssesAttention(nn.Module): self.num_heads = num_heads self.head_size = head_size self.num_kv_heads = num_kv_heads - self.backend = backend_name_to_enum(attn_backend.get_name()) + self.backend = attn_backend.get_enum() self.dtype = dtype @torch.compiler.disable @@ -259,7 +256,7 @@ class LocalAttention(nn.Module): self.num_heads = num_heads self.head_size = head_size self.num_kv_heads = num_kv_heads - self.backend = backend_name_to_enum(attn_backend.get_name()) + self.backend = attn_backend.get_enum() self.dtype = dtype def forward( @@ -336,7 +333,7 @@ class USPAttention(nn.Module): self.num_heads = num_heads self.head_size = head_size self.num_kv_heads = num_kv_heads - self.backend = backend_name_to_enum(attn_backend.get_name()) + self.backend = attn_backend.get_enum() self.dtype = dtype self.causal = causal self.dropout_p = dropout_rate diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index 6145ac0a6..3c6bde541 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -13,30 +13,13 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( 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 import ( + AttentionBackendEnum, + current_platform, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -try: - from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( - SlidingTileAttentionBackend, - ) - - st_attn_available = True -except ImportError: - st_attn_available = False - SlidingTileAttentionBackend = None # type: ignore - -try: - from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( - VideoSparseAttentionBackend, - ) - - vsa_available = True -except ImportError: - vsa_available = False - VideoSparseAttentionBackend = None # type: ignore - logger = init_logger(__name__) @@ -111,7 +94,7 @@ class CausalDMDDenoisingStage(DenoisingStage): ) # STA - if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + if self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN: self.prepare_sta_param(batch, server_args) # Latents and prompts @@ -268,8 +251,8 @@ class CausalDMDDenoisingStage(DenoisingStage): # Attention metadata if needed if ( - vsa_available - and self.attn_backend == VideoSparseAttentionBackend + self.attn_backend.get_enum() + == AttentionBackendEnum.VIDEO_SPARSE_ATTN ): self.attn_metadata_builder_cls = ( self.attn_backend.get_builder_cls() 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 b2e1274ba..7580f7bf4 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -38,13 +38,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, ) - -try: - from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( - FlashAttentionBackend, - ) -except ImportError: - FlashAttentionBackend = None from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import ( configure_sta, @@ -63,41 +56,16 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( 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 import ( + AttentionBackendEnum, + current_platform, +) 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 from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler from sglang.multimodal_gen.utils import dict_to_3d_list, masks_like -try: - from sglang.multimodal_gen.runtime.layers.attention.backends.sliding_tile_attn import ( - SlidingTileAttentionBackend, - ) - - st_attn_available = True -except ImportError: - st_attn_available = False - -try: - from sglang.multimodal_gen.runtime.layers.attention.backends.vmoba import ( - VMOBAAttentionBackend, - ) - from sglang.multimodal_gen.utils import is_vmoba_available - - vmoba_attn_available = is_vmoba_available() -except ImportError: - vmoba_attn_available = False - -try: - from sglang.multimodal_gen.runtime.layers.attention.backends.video_sparse_attn import ( - VideoSparseAttentionBackend, - ) - - vsa_available = True -except ImportError: - vsa_available = False - logger = init_logger(__name__) @@ -563,7 +531,7 @@ class DenoisingStage(PipelineStage): ] # Prepare STA parameters - if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + if self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN: self.prepare_sta_param(batch, server_args) # Get latents and embeddings @@ -732,8 +700,7 @@ class DenoisingStage(PipelineStage): # Save STA mask search results if needed if ( - st_attn_available - and self.attn_backend == SlidingTileAttentionBackend + self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN and server_args.STA_mode == STA_Mode.STA_SEARCHING ): self.save_sta_search_results(batch) @@ -1186,8 +1153,9 @@ class DenoisingStage(PipelineStage): self.attn_metadata_builder_cls = None if self.attn_metadata_builder_cls: self.attn_metadata_builder = self.attn_metadata_builder_cls() - if (st_attn_available and self.attn_backend == SlidingTileAttentionBackend) or ( - vsa_available and self.attn_backend == VideoSparseAttentionBackend + if ( + self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN + or self.attn_backend.get_enum() == AttentionBackendEnum.VIDEO_SPARSE_ATTN ): attn_metadata = self.attn_metadata_builder.build( current_timestep=i, @@ -1197,7 +1165,7 @@ class DenoisingStage(PipelineStage): VSA_sparsity=server_args.VSA_sparsity, device=get_local_torch_device(), ) - elif vmoba_attn_available and self.attn_backend == VMOBAAttentionBackend: + elif self.attn_backend.get_enum() == AttentionBackendEnum.VMOBA_ATTN: moba_params = server_args.moba_config.copy() moba_params.update( { @@ -1207,7 +1175,7 @@ class DenoisingStage(PipelineStage): "device": get_local_torch_device(), } ) - elif self.attn_backend == FlashAttentionBackend: + elif self.attn_backend.get_enum() == AttentionBackendEnum.FA: attn_metadata = self.attn_metadata_builder.build( raw_latent_shape=batch.raw_latent_shape ) diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 3be26231e..19f80f85c 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -23,6 +23,7 @@ logger = init_logger(__name__) class AttentionBackendEnum(enum.Enum): + FA2 = enum.auto() FA = enum.auto() SLIDING_TILE_ATTN = enum.auto() TORCH_SDPA = enum.auto()