[diffusion] feat: add dynamic cuDNN SDPA attention backend (#30090)

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)

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,

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,

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,

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

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: