# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. import warnings import torch from typing import Optional, Tuple try: import flash_attn_interface FLASH_ATTN_3_AVAILABLE = True except ModuleNotFoundError: FLASH_ATTN_3_AVAILABLE = False try: import flash_attn FLASH_ATTN_2_AVAILABLE = True except ModuleNotFoundError: FLASH_ATTN_2_AVAILABLE = False __all__ = [ 'flash_attention', 'attention', ] # --------------------------- # Custom op + fake kernel # --------------------------- from typing import Optional, Sequence # <- add Sequence # ... imports unchanged ... from typing import Optional, Sequence @torch.library.custom_op("wan::flash_attention", mutates_args=()) def _wan_flash_attention_op( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_lens: Optional[torch.Tensor] = None, k_lens: Optional[torch.Tensor] = None, dropout_p: float = 0.0, softmax_scale: Optional[float] = None, q_scale: Optional[float] = None, causal: bool = False, # IMPORTANT: schema-friendly default (None), not a tuple window_size: Optional[Sequence[int]] = None, deterministic: bool = False, dtype: torch.dtype = torch.bfloat16, version: Optional[int] = None, ) -> torch.Tensor: half_dtypes = (torch.float16, torch.bfloat16) assert dtype in half_dtypes assert q.size(-1) <= 256 # normalize window_size to a 2-tuple for FA2 API if window_size is None: ws = (-1, -1) else: ws = tuple(window_size) if len(ws) != 2: raise ValueError(f"window_size must have length 2; got {window_size!r}") b, lq, nheads = q.shape[0], q.shape[1], q.shape[2] lk = k.shape[1] out_dtype = q.dtype def half(x: torch.Tensor) -> torch.Tensor: return x if x.dtype in half_dtypes else x.to(dtype) # --- preprocess (unchanged) --- if q_lens is None: q_flat = half(q.flatten(0, 1)) q_lens = torch.tensor([lq] * b, dtype=torch.int32) else: q_flat = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) if k_lens is None: k_flat = half(k.flatten(0, 1)) v_flat = half(v.flatten(0, 1)) k_lens = torch.tensor([lk] * b, dtype=torch.int32) else: k_flat = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) v_flat = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) q_flat = q_flat.to(v_flat.dtype); k_flat = k_flat.to(v_flat.dtype) if q_scale is not None: q_flat = q_flat * q_scale if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: warnings.warn('Flash attention 3 is not available, use flash attention 2 instead.') if FLASH_ATTN_3_AVAILABLE: ret = flash_attn_interface.flash_attn_varlen_func( q=q_flat, k=k_flat, v=v_flat, cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0, dtype=torch.int32).to(q_flat.device, non_blocking=True), cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(0, dtype=torch.int32).to(k_flat.device, non_blocking=True), seqused_q=None, seqused_k=None, max_seqlen_q=lq, max_seqlen_k=lk, softmax_scale=softmax_scale, causal=causal, deterministic=deterministic, ) out0 = ret[0] if isinstance(ret, (tuple, list)) else ret total_q = b * lq if out0.dim() != 3: raise RuntimeError(f"Unexpected FA3 output rank {out0.dim()} shape={tuple(out0.shape)}") if out0.shape[0] == total_q: out_flat = out0 elif out0.shape[0] == nheads and out0.shape[1] == total_q: out_flat = out0.transpose(0, 1).contiguous() else: raise RuntimeError(f"Unexpected FA3 output shape {tuple(out0.shape)}") out = out_flat.unflatten(0, (b, lq)) elif FLASH_ATTN_2_AVAILABLE: out = flash_attn.flash_attn_varlen_func( q=q_flat, k=k_flat, v=v_flat, cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(0, dtype=torch.int32).to(q_flat.device, non_blocking=True), cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(0, dtype=torch.int32).to(q_flat.device, non_blocking=True), max_seqlen_q=lq, max_seqlen_k=lk, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=causal, window_size=ws, # <- pass 2-tuple deterministic=deterministic, ).unflatten(0, (b, lq)) else: q_s = q.transpose(1, 2).to(dtype) k_s = k.transpose(1, 2).to(dtype) v_s = v.transpose(1, 2).to(dtype) out = torch.nn.functional.scaled_dot_product_attention( q_s, k_s, v_s, attn_mask=None, is_causal=causal, dropout_p=dropout_p ).transpose(1, 2).contiguous() return out.to(out_dtype) @_wan_flash_attention_op.register_fake def _wan_flash_attention_op_fake( q, k, v, q_lens=None, k_lens=None, dropout_p: float = 0.0, softmax_scale=None, q_scale=None, causal: bool = False, window_size: Optional[Sequence[int]] = None, deterministic: bool = False, dtype: torch.dtype = torch.bfloat16, version: Optional[int] = None, ): # Match output shape: (B, Lq, Nq, Dh_v) and keep the SAME fake device as `q` B, Lq, Nq, _ = q.shape Dh_v = v.shape[-1] return q.new_empty((B, Lq, Nq, Dh_v), dtype=q.dtype) # --------------------------- # Public API (unchanged signature) # --------------------------- def flash_attention( q, k, v, q_lens=None, k_lens=None, dropout_p=0., softmax_scale=None, q_scale=None, causal=False, window_size=(-1, -1), deterministic=False, dtype=torch.bfloat16, version=None, ): """ q: [B, Lq, Nq, C1]. k: [B, Lk, Nk, C1]. v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. q_lens: [B]. k_lens: [B]. dropout_p: float. Dropout probability. softmax_scale: float. The scaling of QK^T before applying softmax. causal: bool. Whether to apply causal attention mask. window_size: (left right). If not (-1, -1), apply sliding window local attention. deterministic: bool. If True, slightly slower and uses more memory. dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. """ # Simply delegate to the custom op so Dynamo/AOT treats it as a single node; # our eager kernel inside _wan_flash_attention_op keeps the original behavior. return _wan_flash_attention_op( q, k, v, q_lens=q_lens, k_lens=k_lens, dropout_p=dropout_p, softmax_scale=softmax_scale, q_scale=q_scale, causal=causal, window_size=window_size, deterministic=deterministic, dtype=dtype, version=version, ) def attention( q, k, v, q_lens=None, k_lens=None, dropout_p=0., softmax_scale=None, q_scale=None, causal=False, window_size=(-1, -1), deterministic=False, dtype=torch.bfloat16, fa_version=None, ): if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: return flash_attention( q=q, k=k, v=v, q_lens=q_lens, k_lens=k_lens, dropout_p=dropout_p, softmax_scale=softmax_scale, q_scale=q_scale, causal=causal, window_size=window_size, deterministic=deterministic, dtype=dtype, version=fa_version, ) else: if q_lens is not None or k_lens is not None: warnings.warn( 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' ) q_ = q.transpose(1, 2).to(dtype) k_ = k.transpose(1, 2).to(dtype) v_ = v.transpose(1, 2).to(dtype) out = torch.nn.functional.scaled_dot_product_attention( q_, k_, v_, attn_mask=None, is_causal=causal, dropout_p=dropout_p ) return out.transpose(1, 2).contiguous()