From 969faaa4107369cbbde548c9cfd895695c2e226a Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Wed, 14 Jan 2026 23:31:46 +0800 Subject: [PATCH] [diffusion] fix: revise fa4 backend to support blackwell (#17077) Co-authored-by: luoyuan.luo --- .../layers/attention/backends/flash_attn.py | 240 +++++++++++++++--- 1 file changed, 207 insertions(+), 33 deletions(-) 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 e533b1082..e49c024af 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 @@ -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.")