diff --git a/python/sglang/multimodal_gen/configs/models/dits/base.py b/python/sglang/multimodal_gen/configs/models/dits/base.py index 128238b36..366c816ba 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/base.py +++ b/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, diff --git a/python/sglang/multimodal_gen/configs/models/encoders/base.py b/python/sglang/multimodal_gen/configs/models/encoders/base.py index 1ae63fe92..0c4f86b36 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/base.py +++ b/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, } ) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index ee6cd38b8..021e9db59 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py index 05f30c085..df795e062 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py +++ b/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"]: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 9cce5ac1f..2789ebdf3 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/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, ), ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/flux.py b/python/sglang/multimodal_gen/runtime/models/dits/flux.py index 5bc7dad76..75fa2d3eb 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/flux.py +++ b/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, ), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index f6394e942..ad9e10dd5 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/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, ), ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 97a676891..c8d67ace4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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, }, ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py index 50b15b61b..19286e644 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py index ebac284a7..65bb83280 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py +++ b/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, ), ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py index a151fef86..6b7c6f5f3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 382a928b3..4f1d41b11 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/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") - target_backend = AttentionBackendEnum.FA3 + from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( + 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 - logger.debug(f"Use torch_sdpa as default backend") - else: - target_backend = AttentionBackendEnum.FA3 - logger.debug(f"Use fa3 as default backend") + from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( + set_fa_ver, + ) + + 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" diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 68073b7fb..660f8e1c4 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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() diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index 9eca14ac7..cadf98000 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 1f4367411..43ae2aff6 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/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: