[diffusion] chore: allow all attention backends if not specified (#15530)

python/sglang/multimodal_gen/runtime/layers/attention/selector.py

@@ -91,13 +91,13 @@ def get_attn_backend(
     dtype: torch.dtype,
     supported_attention_backends: set[AttentionBackendEnum] | None = None,
 ) -> type[AttentionBackend]:
-    if supported_attention_backends is not None:
+    if supported_attention_backends is None:
+        be_tuple = tuple()
+    else:
         # Sort the backend names to ensure consistent cache key
         be_tuple = tuple(
             sorted(list(supported_attention_backends), key=lambda b: b.name)
         )
-    else:
-        be_tuple = None
     return _cached_get_attn_backend(head_size, dtype, be_tuple)
 
 
@@ -105,7 +105,7 @@ def get_attn_backend(
 def _cached_get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: tuple[AttentionBackendEnum] | None = None,
+    supported_attention_backends: tuple[AttentionBackendEnum],
 ) -> type[AttentionBackend]:
     # Check whether a particular choice of backend was
     # previously forced.
@@ -115,8 +115,6 @@ def _cached_get_attn_backend(
     from sglang.multimodal_gen.runtime.platforms import current_platform
 
     supported_attention_backends = set(supported_attention_backends)
-    if not supported_attention_backends:
-        raise ValueError("supported_attention_backends is empty")
     selected_backend = None
     backend_by_global_setting: AttentionBackendEnum | None = (
         get_global_forced_attn_backend()
@@ -139,12 +137,12 @@ def _cached_get_attn_backend(
                 )
 
     # get device-specific attn_backend
-    if selected_backend is None:
+    if len(supported_attention_backends) == 0:
+        # all attention backends are allowed
+        pass
+    elif selected_backend is None:
         logger.debug(f"Attention backend not specified")
-    elif (
+    elif selected_backend not in supported_attention_backends:
-        not supported_attention_backends
-        or selected_backend not in supported_attention_backends
-    ):
         supported_attention_backends_str = [
             supported_attention_backend.__str__()
             for supported_attention_backend in supported_attention_backends

python/sglang/multimodal_gen/runtime/models/dits/flux.py

@@ -43,10 +43,7 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     _apply_rotary_emb,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.platforms import (
+from sglang.multimodal_gen.runtime.platforms import current_platform
-    AttentionBackendEnum,
-    current_platform,
-)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
@@ -125,13 +122,6 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
             dropout_rate=0,
             softmax_scale=None,
             causal=False,
-            supported_attention_backends={
-                AttentionBackendEnum.FA,
-                AttentionBackendEnum.AITER,
-                AttentionBackendEnum.TORCH_SDPA,
-                AttentionBackendEnum.SAGE_ATTN,
-                AttentionBackendEnum.SAGE_ATTN_3,
-            },
         )
 
     def forward(

python/sglang/multimodal_gen/runtime/models/dits/flux_2.py

@@ -29,10 +29,7 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     _apply_rotary_emb,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.platforms import (
+from sglang.multimodal_gen.runtime.platforms import current_platform
-    AttentionBackendEnum,
-    current_platform,
-)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
@@ -149,12 +146,6 @@ class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
             dropout_rate=0,
             softmax_scale=None,
             causal=False,
-            supported_attention_backends={
-                AttentionBackendEnum.FA,
-                AttentionBackendEnum.TORCH_SDPA,
-                AttentionBackendEnum.SAGE_ATTN,
-                AttentionBackendEnum.SAGE_ATTN_3,
-            },
         )
 
     def forward(
@@ -283,12 +274,6 @@ class Flux2ParallelSelfAttention(torch.nn.Module, AttentionModuleMixin):
             dropout_rate=0,
             softmax_scale=None,
             causal=False,
-            supported_attention_backends={
-                AttentionBackendEnum.FA,
-                AttentionBackendEnum.TORCH_SDPA,
-                AttentionBackendEnum.SAGE_ATTN,
-                AttentionBackendEnum.SAGE_ATTN_3,
-            },
         )
 
     def forward(

python/sglang/multimodal_gen/runtime/models/dits/zimage.py

@@ -11,10 +11,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.platforms import (
+from sglang.multimodal_gen.runtime.platforms import current_platform
-    AttentionBackendEnum,
-    current_platform,
-)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -125,10 +122,6 @@ class ZImageAttention(nn.Module):
             dropout_rate=0,
             softmax_scale=None,
             causal=False,
-            supported_attention_backends={
-                AttentionBackendEnum.FA,
-                AttentionBackendEnum.TORCH_SDPA,
-            },
         )
 
     def forward(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -64,7 +64,6 @@ 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.interface import AttentionBackendEnum
 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
@@ -133,18 +132,10 @@ class DenoisingStage(PipelineStage):
         self.vae = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
 
+        # TODO(will): hack, should use the actual one in dit
         self.attn_backend = get_attn_backend(
             head_size=attn_head_size,
-            dtype=torch.float16,  # TODO(will): hack
+            dtype=torch.float16,
-            supported_attention_backends={
-                AttentionBackendEnum.SLIDING_TILE_ATTN,
-                AttentionBackendEnum.AITER,
-                AttentionBackendEnum.VIDEO_SPARSE_ATTN,
-                AttentionBackendEnum.VMOBA_ATTN,
-                AttentionBackendEnum.FA,
-                AttentionBackendEnum.TORCH_SDPA,
-                AttentionBackendEnum.SAGE_ATTN_3,
-            },  # hack
         )
 
         # cfg