diffusion: support fa4 in fa backend for blackwell (#13263)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/configs/models/dits/base.py

@@ -20,7 +20,7 @@ class DiTArchConfig(ArchConfig):
         default_factory=lambda: {
             AttentionBackendEnum.SLIDING_TILE_ATTN,
             AttentionBackendEnum.SAGE_ATTN,
-            AttentionBackendEnum.FA3,
+            AttentionBackendEnum.FA,
             AttentionBackendEnum.TORCH_SDPA,
             AttentionBackendEnum.VIDEO_SPARSE_ATTN,
             AttentionBackendEnum.VMOBA_ATTN,

python/sglang/multimodal_gen/configs/models/encoders/base.py

@@ -16,7 +16,7 @@ class EncoderArchConfig(ArchConfig):
     architectures: list[str] = field(default_factory=lambda: [])
     _supported_attention_backends: set[AttentionBackendEnum] = field(
         default_factory=lambda: {
-            AttentionBackendEnum.FA3,
+            AttentionBackendEnum.FA,
             AttentionBackendEnum.TORCH_SDPA,
         }
     )

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py

@@ -28,6 +28,13 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
 
+fa_ver = 3
+
+
+def set_fa_ver(ver: int):
+    global fa_ver
+    fa_ver = ver
+
 
 @dataclass
 class FlashAttentionMetadata:
@@ -128,5 +135,6 @@ class FlashAttentionImpl(AttentionImpl):
             softmax_scale=self.softmax_scale,
             causal=self.causal,
             return_softmax_lse=return_softmax_lse,
+            ver=fa_ver,
         )
         return output

python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py

@@ -27,7 +27,7 @@ class FlashAttention2Backend(AttentionBackend):
 
     @staticmethod
     def get_name() -> str:
-        return "FA3"
+        return "FA"
 
     @staticmethod
     def get_impl_cls() -> type["FlashAttention2Impl"]:

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

@@ -86,7 +86,7 @@ class CausalWanSelfAttention(nn.Module):
             softmax_scale=None,
             causal=False,
             supported_attention_backends=(
-                AttentionBackendEnum.FA3,
+                AttentionBackendEnum.FA,
                 AttentionBackendEnum.TORCH_SDPA,
             ),
         )

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

@@ -156,7 +156,7 @@ class FluxAttention(torch.nn.Module, AttentionModuleMixin):
             softmax_scale=None,
             causal=False,
             supported_attention_backends=(
-                AttentionBackendEnum.FA3,
+                AttentionBackendEnum.FA,
                 AttentionBackendEnum.TORCH_SDPA,
                 AttentionBackendEnum.SAGE_ATTN,
             ),

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

@@ -885,7 +885,7 @@ class IndividualTokenRefinerBlock(nn.Module):
             head_size=hidden_size // num_attention_heads,
             # TODO: remove hardcode; remove STA
             supported_attention_backends=(
-                AttentionBackendEnum.FA3,
+                AttentionBackendEnum.FA,
                 AttentionBackendEnum.TORCH_SDPA,
             ),
         )

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

@@ -289,7 +289,7 @@ class QwenImageCrossAttention(nn.Module):
             softmax_scale=None,
             causal=False,
             supported_attention_backends={
-                AttentionBackendEnum.FA3,
+                AttentionBackendEnum.FA,
                 AttentionBackendEnum.TORCH_SDPA,
             },
         )

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

@@ -157,7 +157,7 @@ class SelfAttention(nn.Module):
         with_qk_norm: bool = True,
         attn_type: str = "torch",
         supported_attention_backends=(
-            AttentionBackendEnum.FA3,
+            AttentionBackendEnum.FA,
             AttentionBackendEnum.TORCH_SDPA,
         ),
     ):
@@ -269,7 +269,7 @@ class CrossAttention(nn.Module):
         bias=False,
         with_qk_norm=True,
         supported_attention_backends=(
-            AttentionBackendEnum.FA3,
+            AttentionBackendEnum.FA,
             AttentionBackendEnum.TORCH_SDPA,
         ),
     ) -> None:

python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py

@@ -137,7 +137,7 @@ class Qwen2_5_VLAttention(nn.Module):
             softmax_scale=self.scaling,
             causal=True,
             supported_attention_backends=(
-                AttentionBackendEnum.FA3,
+                AttentionBackendEnum.FA,
                 AttentionBackendEnum.TORCH_SDPA,
             ),
         )

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

@@ -135,7 +135,7 @@ class DenoisingStage(PipelineStage):
                 AttentionBackendEnum.SLIDING_TILE_ATTN,
                 AttentionBackendEnum.VIDEO_SPARSE_ATTN,
                 AttentionBackendEnum.VMOBA_ATTN,
-                AttentionBackendEnum.FA3,
+                AttentionBackendEnum.FA,
                 AttentionBackendEnum.TORCH_SDPA,
                 AttentionBackendEnum.SAGE_ATTN_THREE,
             },  # hack

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -214,35 +214,45 @@ class CudaPlatformBase(Platform):
         elif selected_backend == AttentionBackendEnum.TORCH_SDPA:
             logger.info("Using Torch SDPA backend.")
             return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
-        elif selected_backend == AttentionBackendEnum.FA3:
+        elif selected_backend in [
+            AttentionBackendEnum.FA,
+        ]:
             if is_blackwell():
-                raise ValueError("The 'fa3' backend is not supported on Blackwell GPUs")
+                from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
-            target_backend = AttentionBackendEnum.FA3
+                    set_fa_ver,
+                )
+
+                set_fa_ver(4)
+            target_backend = AttentionBackendEnum.FA
         elif selected_backend:
             raise ValueError(f"Invalid attention backend for {cls.device_name}")
         else:
             if is_blackwell():
-                target_backend = AttentionBackendEnum.TORCH_SDPA
+                from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (
-                logger.debug(f"Use torch_sdpa as default backend")
+                    set_fa_ver,
-            else:
+                )
-                target_backend = AttentionBackendEnum.FA3
+
-                logger.debug(f"Use fa3 as default backend")
+                set_fa_ver(4)
+            target_backend = AttentionBackendEnum.FA
+            logger.debug(
+                f"Using FlashAttention (FA3 for hopper, FA4 for blackwell) as default backend"
+            )
 
         if not cls.has_device_capability(80):
             logger.info(
-                "Cannot use FlashAttention-2 backend for Volta and Turing " "GPUs."
+                "Cannot use FlashAttention backend for Volta and Turing " "GPUs."
             )
             target_backend = AttentionBackendEnum.TORCH_SDPA
         elif dtype not in (torch.float16, torch.bfloat16):
             logger.info(
-                "Cannot use FlashAttention-2 backend for dtype other than "
+                "Cannot use FlashAttention backend for dtype other than "
                 "torch.float16 or torch.bfloat16."
             )
             target_backend = AttentionBackendEnum.TORCH_SDPA
 
         # FlashAttn is valid for the model, checking if the package is
         # installed.
-        if target_backend == AttentionBackendEnum.FA3:
+        if target_backend == AttentionBackendEnum.FA:
             try:
                 from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import (  # noqa: F401
                     FlashAttentionBackend,
@@ -251,13 +261,13 @@ class CudaPlatformBase(Platform):
                 supported_sizes = FlashAttentionBackend.get_supported_head_sizes()
                 if head_size not in supported_sizes:
                     logger.info(
-                        "Cannot use FlashAttention-2 backend for head size %d.",
+                        "Cannot use FlashAttention backend for head size %d.",
                         head_size,
                     )
                     target_backend = AttentionBackendEnum.TORCH_SDPA
             except ImportError:
                 logger.info(
-                    "Cannot use FlashAttention-2 backend because the "
+                    "Cannot use FlashAttention backend because the "
                     "flash_attn package is not found. "
                     "Make sure that flash_attn was built and installed "
                     "(on by default)."
@@ -269,7 +279,7 @@ class CudaPlatformBase(Platform):
 
             return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
 
-        logger.info("Using fa3 backend.")
+        logger.info("Using FlashAttention (FA3 for hopper, FA4 for blackwell) backend.")
 
         return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend"
 

python/sglang/multimodal_gen/runtime/platforms/interface.py

@@ -23,7 +23,7 @@ logger = init_logger(__name__)
 
 
 class AttentionBackendEnum(enum.Enum):
-    FA3 = enum.auto()
+    FA = enum.auto()
     SLIDING_TILE_ATTN = enum.auto()
     TORCH_SDPA = enum.auto()
     SAGE_ATTN = enum.auto()

python/sglang/multimodal_gen/runtime/platforms/rocm.py

@@ -77,7 +77,7 @@ class RocmPlatform(Platform):
             logger.info("Using Torch SDPA backend.")
             return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend"
 
-        elif selected_backend in (AttentionBackendEnum.FA3, None):
+        elif selected_backend in (AttentionBackendEnum.FA, None):
             pass
 
         elif selected_backend in (
@@ -92,7 +92,7 @@ class RocmPlatform(Platform):
                 f"Invalid attention backend for {cls.device_name}: {selected_backend}"
             )
 
-        target_backend = AttentionBackendEnum.FA3
+        target_backend = AttentionBackendEnum.FA
         if dtype not in (torch.float16, torch.bfloat16):
             logger.info(
                 "Cannot use FlashAttention backend for dtype other than "
@@ -100,7 +100,7 @@ class RocmPlatform(Platform):
             )
             target_backend = AttentionBackendEnum.TORCH_SDPA
 
-        if target_backend == AttentionBackendEnum.FA3:
+        if target_backend == AttentionBackendEnum.FA:
             try:
                 import flash_attn  # noqa: F401
 

python/sglang/multimodal_gen/runtime/server_args.py

@@ -336,6 +336,9 @@ class ServerArgs:
 
     def __post_init__(self):
         # Add randomization to avoid race condition when multiple servers start simultaneously
+        if self.attention_backend in ["fa3", "fa4"]:
+            self.attention_backend = "fa"
+
         initial_scheduler_port = self.scheduler_port + random.randint(0, 100)
         self.scheduler_port = self.settle_port(initial_scheduler_port)
         # TODO: remove hard code
@@ -382,7 +385,7 @@ class ServerArgs:
             "--attention-backend",
             type=str,
             default=None,
-            choices=[e.name.lower() for e in AttentionBackendEnum],
+            choices=[e.name.lower() for e in AttentionBackendEnum] + ["fa3", "fa4"],
             help="The attention backend to use. If not specified, the backend is automatically selected based on hardware and installed packages.",
         )
 
@@ -843,14 +846,14 @@ class ServerArgs:
             )
 
         if self.ring_degree > 1:
-            if self.attention_backend != None and self.attention_backend != "fa3":
+            if self.attention_backend != None and self.attention_backend != "fa":
                 raise ValueError(
-                    "Ring Attention is only supported for fa3 backend for now"
+                    "Ring Attention is only supported for flash attention backend for now"
                 )
             else:
-                self.attention_backend = "fa3"
+                self.attention_backend = "fa"
                 logger.info(
-                    "Ring Attention is currently only supported for fa3, attention_backend has been automatically set to fa3"
+                    "Ring Attention is currently only supported for flash attention, attention_backend has been automatically set to flash attention"
                 )
 
         if self.sp_degree == -1: