From a24906a0917056116b66adc11a3dffaac9b9518b Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 28 Jul 2026 19:50:24 +0800 Subject: [PATCH] [diffusion] feat: add dynamic cuDNN SDPA attention backend (#30090) --- .../runtime/layers/attention/backends/sdpa.py | 105 +++++++++++++++++- .../runtime/layers/attention/selector.py | 21 +++- .../memory_managers/component_manager.py | 13 ++- .../multimodal_gen/runtime/platforms/cuda.py | 16 +++ .../runtime/platforms/interface.py | 2 + .../runtime/server_args/server_args.py | 2 + 6 files changed, 149 insertions(+), 10 deletions(-) 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 9345785ed..c117f43f7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py @@ -63,6 +63,11 @@ class SDPAImpl(AttentionImpl): self.dropout = extra_impl_args.get("dropout_p", 0.0) self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False)) + def _sdpa_context(self, query: torch.Tensor): + if self.allow_cudnn_sdp and query.device.type == "cuda": + return sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS) + return nullcontext() + def forward( self, query: torch.Tensor, @@ -82,14 +87,102 @@ class SDPAImpl(AttentionImpl): } if query.shape[1] != key.shape[1]: attn_kwargs["enable_gqa"] = True - sdpa_context = ( - sdpa_kernel(_PYTORCH_DEFAULT_CUDA_SDP_BACKENDS) - if self.allow_cudnn_sdp and query.device.type == "cuda" - else nullcontext() - ) - with sdpa_context: + with self._sdpa_context(query): output = torch.nn.functional.scaled_dot_product_attention( query, key, value, **attn_kwargs ) output = output.transpose(1, 2) return output + + +class CudnnSDPABackend(SDPABackend): + @staticmethod + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.TORCH_CUDNN_SDPA + + @staticmethod + def get_impl_cls() -> type["CudnnSDPAImpl"]: + return CudnnSDPAImpl + + +class CudnnSDPAImpl(SDPAImpl): + def _sdpa_context(self, query: torch.Tensor): + if query.device.type == "cuda": + return sdpa_kernel(SDPBackend.CUDNN_ATTENTION) + return nullcontext() + + +class DynamicCudnnSDPABackend(SDPABackend): + @staticmethod + def get_enum() -> AttentionBackendEnum: + return AttentionBackendEnum.DYNAMIC_CUDNN_SDPA + + @staticmethod + def get_impl_cls() -> type["DynamicCudnnSDPAImpl"]: + return DynamicCudnnSDPAImpl + + +class DynamicCudnnSDPAImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool, + softmax_scale: float, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( + FlashAttentionImpl, + set_fa_ver, + ) + + self.causal = causal + self.head_size = head_size + if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10: + set_fa_ver(4) + self.cudnn_impl = CudnnSDPAImpl( + num_heads=num_heads, + head_size=head_size, + causal=causal, + softmax_scale=softmax_scale, + num_kv_heads=num_kv_heads, + prefix=f"{prefix}.cudnn", + **extra_impl_args, + ) + self.fa_impl = FlashAttentionImpl( + num_heads=num_heads, + head_size=head_size, + causal=causal, + softmax_scale=softmax_scale, + num_kv_heads=num_kv_heads, + prefix=f"{prefix}.fa", + **extra_impl_args, + ) + + def _use_cudnn_sdpa( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> bool: + if self.causal: + return False + if query.device.type != "cuda": + return False + if query.dtype not in (torch.float16, torch.bfloat16): + return False + if query.shape[2] != key.shape[2]: + return False + if query.shape[1] != key.shape[1]: + return False + return query.shape[-1] == 64 and query.shape[1] == 1024 and query.shape[0] >= 4 + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + if self._use_cudnn_sdpa(query, key, value): + return self.cudnn_impl.forward(query, key, value, attn_metadata) + return self.fa_impl.forward(query, key, value, attn_metadata) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index 07f61baf6..aeea8ed2d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -211,9 +211,8 @@ 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 selected_backend not in 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__() @@ -236,6 +235,22 @@ def _cached_get_attn_backend( return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls)) +def _is_backend_supported( + selected_backend: AttentionBackendEnum, + supported_attention_backends: set[AttentionBackendEnum], +) -> bool: + if selected_backend in supported_attention_backends: + return True + if selected_backend == AttentionBackendEnum.TORCH_CUDNN_SDPA: + return AttentionBackendEnum.TORCH_SDPA in supported_attention_backends + if selected_backend == AttentionBackendEnum.DYNAMIC_CUDNN_SDPA: + return ( + AttentionBackendEnum.FA in supported_attention_backends + and AttentionBackendEnum.TORCH_SDPA in supported_attention_backends + ) + return False + + @contextmanager def component_attn_backend_context_manager( attn_backend: AttentionBackendEnum | None, diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py index 57d89315c..63eec002f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py @@ -186,7 +186,10 @@ class ComponentResidencyManager: ) -> None: """A hook called before processing an actual request""" self.refresh_server_args(server_args) - self.state = ResidencyState(stages=stages, batch_is_warmup=batch.is_warmup) + self.state = ResidencyState( + stages=stages, + batch_is_warmup=self._is_warmup_batch(batch), + ) self._active_use = None self._active_use_module = None self._disable_active_nvtx() @@ -201,6 +204,14 @@ class ComponentResidencyManager: use for uses in self._stage_uses_by_index for use in uses ) + @staticmethod + def _is_warmup_batch(batch: ResidencyBatch | list[ResidencyBatch]) -> bool: + if isinstance(batch, list): + return bool(batch) and all( + getattr(item, "is_warmup", False) for item in batch + ) + return batch.is_warmup + def before_stage( self, stage: ComponentResidencyStage, diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 0cb5d4a7a..b32dfe30a 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -30,6 +30,10 @@ logger = init_logger(__name__) _SDPA_BACKEND_CLS_STR = ( "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" ) +_CUDNN_SDPA_BACKEND_CLS_STR = ( + "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.CudnnSDPABackend" +) +_DYNAMIC_CUDNN_SDPA_BACKEND_CLS_STR = "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.DynamicCudnnSDPABackend" _P = ParamSpec("_P") _R = TypeVar("_R") @@ -100,6 +104,16 @@ class _TorchSDPAAttentionBackendResolver(_DirectCudaAttentionBackendResolver): backend_cls_str = _SDPA_BACKEND_CLS_STR +class _TorchCudnnSDPAAttentionBackendResolver(_DirectCudaAttentionBackendResolver): + backend = AttentionBackendEnum.TORCH_CUDNN_SDPA + backend_cls_str = _CUDNN_SDPA_BACKEND_CLS_STR + + +class _DynamicCudnnSDPAAttentionBackendResolver(_DirectCudaAttentionBackendResolver): + backend = AttentionBackendEnum.DYNAMIC_CUDNN_SDPA + backend_cls_str = _DYNAMIC_CUDNN_SDPA_BACKEND_CLS_STR + + class _SparseLinearAttentionBackendResolver(_DirectCudaAttentionBackendResolver): backend = AttentionBackendEnum.SLA_ATTN backend_cls_str = "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SparseLinearAttentionBackend" @@ -270,6 +284,8 @@ _CUDA_ATTENTION_BACKEND_RESOLVERS = { for resolver in ( _AITerAttentionBackendResolver, _TorchSDPAAttentionBackendResolver, + _TorchCudnnSDPAAttentionBackendResolver, + _DynamicCudnnSDPAAttentionBackendResolver, _SparseLinearAttentionBackendResolver, _SageSparseLinearAttentionBackendResolver, _SlidingTileAttentionBackendResolver, diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index ca136d39c..12f1f2215 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/python/sglang/multimodal_gen/runtime/platforms/interface.py @@ -29,6 +29,8 @@ class AttentionBackendEnum(enum.Enum): FA = enum.auto() SLIDING_TILE_ATTN = enum.auto() TORCH_SDPA = enum.auto() + TORCH_CUDNN_SDPA = enum.auto() + DYNAMIC_CUDNN_SDPA = enum.auto() SAGE_ATTN = enum.auto() SAGE_ATTN_3 = enum.auto() VIDEO_SPARSE_ATTN = enum.auto() diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index a722333c0..8649395fa 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -802,6 +802,8 @@ class ServerArgs(DisaggServerArgsMixin): normalized = backend.strip().lower() if normalized in ("fa3", "fa4"): normalized = "fa" + elif normalized == "cudnn_sdpa": + normalized = "torch_cudnn_sdpa" try: return AttentionBackendEnum[normalized.upper()].name.lower() except KeyError: