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

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

Browse Source 
Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
This commit is contained in:
R0CKSTAR
2025-12-25 11:51:40 +08:00
committed by GitHub
parent b311c43d13
commit 8865424f14
14 changed files with 53 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/aiter.py
View File

@@ -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"]:

4
python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py
View File

@@ -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

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
View File

@@ -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"]:

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py
View File

@@ -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"]:

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn.py
View File

@@ -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"]:

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/sage_attn3.py
View File

@@ -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"]:

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py
View File

@@ -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"]:

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/sliding_tile_attn.py
View File

@@ -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"]:

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/video_sparse_attn.py
View File

@@ -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"]:

5
python/sglang/multimodal_gen/runtime/layers/attention/backends/vmoba.py
View File

@@ -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"]:

11
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
View File

@@ -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

31
python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py
View File

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

54
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

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

1
python/sglang/multimodal_gen/runtime/platforms/interface.py
View File

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