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

[diffusion] fix: revise fa4 backend to support blackwell (#17077)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-01-14 23:31:46 +08:00
committed by GitHub
parent 48c2aca9ba
commit 969faaa410
1 changed files with 207 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,5 +1,4 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass
from functools import lru_cache
@@ -21,11 +20,18 @@ except ImportError as e:
raise e
def maybe_contiguous(x):
def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
return x.contiguous() if x is not None and x.stride(-1) != 1 else x
def flash_attn_varlen_func_fake(
# -----------------------------
# Fake implementations for schema / tracing
# custom op schema requires FIXED return structure.
# We provide TWO ops:
# 1) out-only op: always returns Tensor
# 2) out+lse op: always returns Tuple[Tensor, Tensor]
# -----------------------------
def flash_attn_varlen_func_fake_out(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -51,6 +57,66 @@ def flash_attn_varlen_func_fake(
return_softmax_lse: bool = False,
sinks: Optional[torch.Tensor] = None,
ver: int = 4,
) -> torch.Tensor:
assert ver == 4, "only support flash attention v4"
q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
num_head, head_dim = q.shape[-2:]
if cu_seqlens_q is None:
batch_size, seqlen_q = q.shape[:2]
else:
batch_size = cu_seqlens_q.shape[0] - 1
seqlen_q = None
head_dim_v = v.shape[-1]
if cu_seqlens_q is not None:
assert cu_seqlens_q.shape == (
batch_size + 1,
), "cu_seqlens_q must have shape (batch_size + 1,)"
assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
assert q.dtype in [
torch.float16,
torch.bfloat16,
], "inputs must be float16 or bfloat16"
assert q.dtype == k.dtype == v.dtype, "inputs must have the same dtype"
assert head_dim <= 256, "head_dim must be less than or equal to 256"
alignment = 16 // q.element_size()
assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}"
q_batch_seqlen_shape = (
(batch_size, seqlen_q) if cu_seqlens_q is None else (q.shape[0],)
)
out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
return out
def flash_attn_varlen_func_fake_out_lse(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: Optional[torch.Tensor] = None,
cu_seqlens_k: Optional[torch.Tensor] = None,
max_seqlen_q: Optional[int] = None,
max_seqlen_k: Optional[int] = None,
seqused_q: Optional[torch.Tensor] = None,
seqused_k: Optional[torch.Tensor] = None,
page_table: Optional[torch.Tensor] = None,
softmax_scale: Optional[float] = None,
causal: bool = False,
qv: Optional[torch.Tensor] = None,
q_descale: Optional[torch.Tensor] = None,
k_descale: Optional[torch.Tensor] = None,
v_descale: Optional[torch.Tensor] = None,
window_size: Optional[List[int]] = None,
attention_chunk: int = 0,
softcap: float = 0.0,
num_splits: int = 1,
pack_gqa: Optional[bool] = None,
sm_margin: int = 0,
return_softmax_lse: bool = True,
sinks: Optional[torch.Tensor] = None,
ver: int = 4,
) -> Tuple[torch.Tensor, torch.Tensor]:
assert ver == 4, "only support flash attention v4"
q, k, v = [maybe_contiguous(t) for t in (q, k, v)]
@@ -63,12 +129,14 @@ def flash_attn_varlen_func_fake(
seqlen_q = None
total_q = q.shape[0]
head_dim_v = v.shape[-1]
if cu_seqlens_q is not None:
assert cu_seqlens_q.shape == (
batch_size + 1,
), "cu_seqlens_q must have shape (batch_size + 1,)"
assert cu_seqlens_q.dtype == torch.int32, "cu_seqlens_q must be int32"
assert cu_seqlens_q.stride(0) == 1, "cu_seqlens_q must be contiguous"
assert q.dtype in [
torch.float16,
torch.bfloat16,
@@ -86,12 +154,18 @@ def flash_attn_varlen_func_fake(
if cu_seqlens_q is None
else (num_head, total_q)
)
out = q.new_empty(*q_batch_seqlen_shape, num_head, head_dim_v)
lse = q.new_empty(lse_shape, dtype=torch.float32) if return_softmax_lse else None
return (out, lse) if return_softmax_lse else out
lse = q.new_empty(lse_shape, dtype=torch.float32)
return out, lse
@register_custom_op(fake_impl=flash_attn_varlen_func_fake)
# -----------------------------
# Registered custom ops
# NOTE: fixed return schemas to avoid:
# "Object of type 'Tensor' is not an instance of 'sequence'"
# -----------------------------
@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out)
def flash_attn_varlen_func_op(
q: torch.Tensor,
k: torch.Tensor,
@@ -118,9 +192,14 @@ def flash_attn_varlen_func_op(
return_softmax_lse: bool = False,
sinks: Optional[torch.Tensor] = None,
ver: int = 4,
) -> Tuple[torch.Tensor, torch.Tensor]:
) -> torch.Tensor:
if window_size is None:
window_size = [-1, -1]
if return_softmax_lse:
raise ValueError(
"flash_attn_varlen_func_op is out-only op; return_softmax_lse must be False. "
"Use flash_attn_varlen_func_op_lse for (out, lse)."
)
return flash_attn_func(
q,
k,
@@ -144,7 +223,71 @@ def flash_attn_varlen_func_op(
num_splits=num_splits,
pack_gqa=pack_gqa,
sm_margin=sm_margin,
return_softmax_lse=return_softmax_lse,
return_softmax_lse=False,
sinks=sinks,
ver=ver,
)
@register_custom_op(fake_impl=flash_attn_varlen_func_fake_out_lse)
def flash_attn_varlen_func_op_lse(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: Optional[torch.Tensor] = None,
cu_seqlens_k: Optional[torch.Tensor] = None,
max_seqlen_q: Optional[int] = None,
max_seqlen_k: Optional[int] = None,
seqused_q: Optional[torch.Tensor] = None,
seqused_k: Optional[torch.Tensor] = None,
page_table: Optional[torch.Tensor] = None,
softmax_scale: Optional[float] = None,
causal: bool = False,
qv: Optional[torch.Tensor] = None,
q_descale: Optional[torch.Tensor] = None,
k_descale: Optional[torch.Tensor] = None,
v_descale: Optional[torch.Tensor] = None,
window_size: Optional[List[int]] = None,
attention_chunk: int = 0,
softcap: float = 0.0,
num_splits: int = 1,
pack_gqa: Optional[bool] = None,
sm_margin: int = 0,
return_softmax_lse: bool = True,
sinks: Optional[torch.Tensor] = None,
ver: int = 4,
) -> Tuple[torch.Tensor, torch.Tensor]:
if window_size is None:
window_size = [-1, -1]
if not return_softmax_lse:
raise ValueError(
"flash_attn_varlen_func_op_lse is out+lse op; return_softmax_lse must be True. "
"Use flash_attn_varlen_func_op for out-only."
)
return flash_attn_func(
q,
k,
v,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
seqused_q=seqused_q,
seqused_k=seqused_k,
page_table=page_table,
softmax_scale=softmax_scale,
causal=causal,
qv=qv,
q_descale=q_descale,
k_descale=k_descale,
v_descale=v_descale,
window_size=tuple(window_size),
attention_chunk=attention_chunk,
softcap=softcap,
num_splits=num_splits,
pack_gqa=pack_gqa,
sm_margin=sm_margin,
return_softmax_lse=True,
sinks=sinks,
ver=ver,
)
@@ -215,7 +358,7 @@ def _should_use_upstream_flash_attention(
return True
def set_fa_ver(ver: int):
def set_fa_ver(ver: int) -> None:
global fa_ver
fa_ver = ver
@@ -234,11 +377,10 @@ class FlashAttentionMetadata:
class FlashAttentionMetadataBuilder(AttentionMetadataBuilder):
def __init__(self):
def __init__(self) -> None:
pass
def prepare(self):
def prepare(self) -> None:
pass
def build( # type: ignore
@@ -275,7 +417,6 @@ class FlashAttentionBackend(AttentionBackend):
class FlashAttentionImpl(AttentionImpl):
def __init__(
self,
num_heads: int,
@@ -318,9 +459,11 @@ class FlashAttentionImpl(AttentionImpl):
else:
max_seqlen_q = query.shape[1]
max_seqlen_k = key.shape[1]
q_shape = tuple(query.shape)
k_shape = tuple(key.shape)
v_shape = tuple(value.shape)
use_upstream = _should_use_upstream_flash_attention(
flash_attn_varlen_func_upstream is not None,
self._upstream_heads_ok,
@@ -347,28 +490,59 @@ class FlashAttentionImpl(AttentionImpl):
return_attn_probs=return_softmax_lse,
)
if return_softmax_lse:
out, softmax_lse = out
return out.reshape(bsz, seqlen, nheads_q, -1), softmax_lse
out_tensor, softmax_lse = out
return out_tensor.reshape(bsz, seqlen, nheads_q, -1), softmax_lse
return out.reshape(bsz, seqlen, nheads_q, d)
# FA version selection:
# - fa_ver == 3: call python function (can return Tensor or (Tensor, Tensor) depending on flag)
# - fa_ver == 4: call custom ops with FIXED return schema
if fa_ver == 3:
flash_attn_op = flash_attn_func
elif fa_ver == 4:
flash_attn_op = flash_attn_varlen_func_op
else:
raise ValueError(f"flash attention version {fa_ver} is not supported.")
output = flash_attn_op(
q=query,
k=key,
v=value,
cu_seqlens_q=None,
cu_seqlens_k=None,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.softmax_scale,
causal=self.causal,
return_softmax_lse=return_softmax_lse,
ver=fa_ver,
)
return output
output = flash_attn_op(
q=query, # type: ignore[no-untyped-call]
k=key,
v=value,
cu_seqlens_q=None,
cu_seqlens_k=None,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.softmax_scale,
causal=self.causal,
return_softmax_lse=return_softmax_lse,
ver=fa_ver,
)
return output
if fa_ver == 4:
if return_softmax_lse:
out_tensor, softmax_lse = flash_attn_varlen_func_op_lse(
q=query,
k=key,
v=value,
cu_seqlens_q=None,
cu_seqlens_k=None,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.softmax_scale,
causal=self.causal,
return_softmax_lse=True,
ver=fa_ver,
)
return out_tensor, softmax_lse
out_tensor = flash_attn_varlen_func_op(
q=query,
k=key,
v=value,
cu_seqlens_q=None,
cu_seqlens_k=None,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
softmax_scale=self.softmax_scale,
causal=self.causal,
return_softmax_lse=False,
ver=fa_ver,
)
return out_tensor
raise ValueError(f"flash attention version {fa_ver} is not supported.")