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

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

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
+    lse = q.new_empty(lse_shape, dtype=torch.float32)
-    return (out, lse) if return_softmax_lse else out
+    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) -> None:
-    def __init__(self):
         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
+                out_tensor, softmax_lse = out
-                return out.reshape(bsz, seqlen, nheads_q, -1), softmax_lse
+                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:
+            output = flash_attn_op(
-            flash_attn_op = flash_attn_varlen_func_op
+                q=query,
-        else:
+                k=key,
-            raise ValueError(f"flash attention version {fa_ver} is not supported.")
+                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(
+        if fa_ver == 4:
-            q=query,  # type: ignore[no-untyped-call]
+            if return_softmax_lse:
-            k=key,
+                out_tensor, softmax_lse = flash_attn_varlen_func_op_lse(
-            v=value,
+                    q=query,
-            cu_seqlens_q=None,
+                    k=key,
-            cu_seqlens_k=None,
+                    v=value,
-            max_seqlen_q=max_seqlen_q,
+                    cu_seqlens_q=None,
-            max_seqlen_k=max_seqlen_k,
+                    cu_seqlens_k=None,
-            softmax_scale=self.softmax_scale,
+                    max_seqlen_q=max_seqlen_q,
-            causal=self.causal,
+                    max_seqlen_k=max_seqlen_k,
-            return_softmax_lse=return_softmax_lse,
+                    softmax_scale=self.softmax_scale,
-            ver=fa_ver,
+                    causal=self.causal,
-        )
+                    return_softmax_lse=True,
-        return output
+                    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.")