[diffusion] refactor: refactor attention backend checking to use backend enum (#15555)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "AITER"
+        return AttentionBackendEnum.AITER
 
     @staticmethod
     def get_impl_cls() -> type["AITerImpl"]:

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

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "FLASH_ATTN"
+        return AttentionBackendEnum.FA
 
     @staticmethod
     def get_impl_cls() -> type["FlashAttentionImpl"]:

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "FA"
+        return AttentionBackendEnum.FA2
 
     @staticmethod
     def get_impl_cls() -> type["FlashAttention2Impl"]:

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "SAGE_ATTN"
+        return AttentionBackendEnum.SAGE_ATTN
 
     @staticmethod
     def get_impl_cls() -> type["SageAttentionImpl"]:

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "SAGE_ATTN_3"
+        return AttentionBackendEnum.SAGE_ATTN_3
 
     @staticmethod
     def get_impl_cls() -> type["SageAttention3Impl"]:

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "SDPA"
+        return AttentionBackendEnum.TORCH_SDPA
 
     @staticmethod
     def get_impl_cls() -> type["SDPAImpl"]:

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "SLIDING_TILE_ATTN"
+        return AttentionBackendEnum.SLIDING_TILE_ATTN
 
     @staticmethod
     def get_impl_cls() -> type["SlidingTileAttentionImpl"]:

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "VIDEO_SPARSE_ATTN"
+        return AttentionBackendEnum.VIDEO_SPARSE_ATTN
 
     @staticmethod
     def get_impl_cls() -> type["VideoSparseAttentionImpl"]:

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:
+    def get_enum() -> AttentionBackendEnum:
-        return "VMOBA_ATTN"
+        return AttentionBackendEnum.VMOBA_ATTN
 
     @staticmethod
     def get_impl_cls() -> type["VMOBAAttentionImpl"]:

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 (
+from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
-    backend_name_to_enum,
-    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

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
+                        self.attn_backend.get_enum()
-                        and self.attn_backend == VideoSparseAttentionBackend
+                        == AttentionBackendEnum.VIDEO_SPARSE_ATTN
                     ):
                         self.attn_metadata_builder_cls = (
                             self.attn_backend.get_builder_cls()

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
+            self.attn_backend.get_enum() == AttentionBackendEnum.SLIDING_TILE_ATTN
-            and self.attn_backend == SlidingTileAttentionBackend
             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 (
+        if (
-            vsa_available and self.attn_backend == VideoSparseAttentionBackend
+            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
             )

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()