chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Yuhao Yang
2025-11-16 21:02:45 +08:00
committed by GitHub
parent 191f5c7795
commit 6afe396399
15 changed files with 55 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/multimodal_gen/configs/models/dits/base.py
View File

@@ -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,

2
python/sglang/multimodal_gen/configs/models/encoders/base.py
View File

@@ -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,
}
)

8
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
View File

@@ -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

2
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn_2.py
View File

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

2
python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py
View File

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

2
python/sglang/multimodal_gen/runtime/models/dits/flux.py
View File

@@ -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,
),

2
python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py
View File

@@ -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,
),
)

2
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

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

4
python/sglang/multimodal_gen/runtime/models/dits/stepvideo.py
View File

@@ -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:

2
python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py
View File

@@ -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,
),
)

2
python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py
View File

@@ -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

38
python/sglang/multimodal_gen/runtime/platforms/cuda.py
View File

@@ -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"

2
python/sglang/multimodal_gen/runtime/platforms/interface.py
View File

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

6
python/sglang/multimodal_gen/runtime/platforms/rocm.py
View File

@@ -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

13
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -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: