Recover attention_dispatch.py with its origin impl, later would special commit for fa3 compatibility.

Author: JerryWu-code

src/diffusers/models/attention_dispatch.py

@@ -18,17 +18,7 @@
 import math
 from dataclasses import dataclass
 from enum import Enum
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    Tuple,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 import torch
 
@@ -78,10 +68,7 @@
 
 if _CAN_USE_FLASH_ATTN:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
-    from flash_attn.flash_attn_interface import (
-        _wrapped_flash_attn_backward,
-        _wrapped_flash_attn_forward,
-    )
+    from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward
 else:
     flash_attn_func = None
     flash_attn_varlen_func = None
@@ -90,13 +77,11 @@
 
 
 if _CAN_USE_FLASH_ATTN_3:
-    from flash_attn_interface import _flash_attn_forward as flash_attn_3_forward
     from flash_attn_interface import flash_attn_func as flash_attn_3_func
     from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
 else:
     flash_attn_3_func = None
     flash_attn_3_varlen_func = None
-    flash_attn_3_forward = None
 
 if _CAN_USE_AITER_ATTN:
     from aiter import flash_attn_func as aiter_flash_attn_func
@@ -135,9 +120,7 @@
 
 
 if _CAN_USE_XLA_ATTN:
-    from torch_xla.experimental.custom_kernel import (
-        flash_attention as xla_flash_attention,
-    )
+    from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention
 else:
     xla_flash_attention = None
 
@@ -280,17 +263,13 @@ class _HubKernelConfig:
 _HUB_KERNELS_REGISTRY: Dict["AttentionBackendName", _HubKernelConfig] = {
     # TODO: temporary revision for now. Remove when merged upstream into `main`.
     AttentionBackendName._FLASH_3_HUB: _HubKernelConfig(
-        repo_id="kernels-community/flash-attn3",
-        function_attr="flash_attn_func",
-        revision="fake-ops-return-probs",
+        repo_id="kernels-community/flash-attn3", function_attr="flash_attn_func", revision="fake-ops-return-probs"
     )
 }
 
 
 @contextlib.contextmanager
-def attention_backend(
-    backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE,
-):
+def attention_backend(backend: Union[str, AttentionBackendName] = AttentionBackendName.NATIVE):
     """
     Context manager to set the active attention backend.
     """
@@ -435,10 +414,7 @@ def _check_attention_backend_requirements(backend: AttentionBackendName) -> None
                 f"Flash Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `flash-attn>={_REQUIRED_FLASH_VERSION}`."
             )
 
-    elif backend in [
-        AttentionBackendName._FLASH_3,
-        AttentionBackendName._FLASH_VARLEN_3,
-    ]:
+    elif backend in [AttentionBackendName._FLASH_3, AttentionBackendName._FLASH_VARLEN_3]:
         if not _CAN_USE_FLASH_ATTN_3:
             raise RuntimeError(
                 f"Flash Attention 3 backend '{backend.value}' is not usable because of missing package or the version is too old. Please build FA3 beta release from source."
@@ -510,11 +486,7 @@ def _prepare_for_flash_attn_or_sage_varlen_without_mask(
     cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0)
     max_seqlen_q = seqlens_q.max().item()
     max_seqlen_k = seqlens_k.max().item()
-    return (
-        (seqlens_q, seqlens_k),
-        (cu_seqlens_q, cu_seqlens_k),
-        (max_seqlen_q, max_seqlen_k),
-    )
+    return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k)
 
 
 def _prepare_for_flash_attn_or_sage_varlen_with_mask(
@@ -531,11 +503,7 @@ def _prepare_for_flash_attn_or_sage_varlen_with_mask(
     cu_seqlens_k[1:] = torch.cumsum(seqlens_k, dim=0)
     max_seqlen_q = seqlens_q.max().item()
     max_seqlen_k = seqlens_k.max().item()
-    return (
-        (seqlens_q, seqlens_k),
-        (cu_seqlens_q, cu_seqlens_k),
-        (max_seqlen_q, max_seqlen_k),
-    )
+    return (seqlens_q, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k)
 
 
 def _prepare_for_flash_attn_or_sage_varlen(
@@ -653,42 +621,22 @@ def _wrapped_flash_attn_3(
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     # Hardcoded for now because pytorch does not support tuple/int type hints
     window_size = (-1, -1)
-    max_seqlen_q = q.shape[2]
-    max_seqlen_k = k.shape[2]
-
-    out, lse, *_ = flash_attn_3_forward(
+    out, lse, *_ = flash_attn_3_func(
         q=q,
         k=k,
         v=v,
-        k_new=None,
-        v_new=None,
+        softmax_scale=softmax_scale,
+        causal=causal,
         qv=qv,
-        out=None,
-        cu_seqlens_q=None,
-        cu_seqlens_k=None,
-        cu_seqlens_k_new=None,
-        seqused_q=None,
-        seqused_k=None,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
-        page_table=None,
-        kv_batch_idx=None,
-        leftpad_k=None,
-        rotary_cos=None,
-        rotary_sin=None,
-        seqlens_rotary=None,
         q_descale=q_descale,
         k_descale=k_descale,
         v_descale=v_descale,
-        softmax_scale=softmax_scale,
-        causal=causal,
         window_size=window_size,
         attention_chunk=attention_chunk,
         softcap=softcap,
-        rotary_interleaved=True,
-        scheduler_metadata=None,
         num_splits=num_splits,
         pack_gqa=pack_gqa,
+        deterministic=deterministic,
         sm_margin=sm_margin,
     )
     lse = lse.permute(0, 2, 1)
@@ -794,10 +742,7 @@ def _native_attention_backward_op(
 
     grad_out_t = grad_out.permute(0, 2, 1, 3)
     grad_query_t, grad_key_t, grad_value_t = torch.autograd.grad(
-        outputs=out,
-        inputs=[query_t, key_t, value_t],
-        grad_outputs=grad_out_t,
-        retain_graph=False,
+        outputs=out, inputs=[query_t, key_t, value_t], grad_outputs=grad_out_t, retain_graph=False
     )
 
     grad_query = grad_query_t.permute(0, 2, 1, 3)
@@ -836,26 +781,18 @@ def _cudnn_attention_forward_op(
     value = value.transpose(1, 2).contiguous()
     tensors_to_save += (query, key, value)
 
-    (
-        out,
-        lse,
-        cum_seq_q,
-        cum_seq_k,
-        max_q,
-        max_k,
-        philox_seed,
-        philox_offset,
-        debug_attn_mask,
-    ) = torch.ops.aten._scaled_dot_product_cudnn_attention(
-        query=query,
-        key=key,
-        value=value,
-        attn_bias=attn_mask,
-        compute_log_sumexp=return_lse,
-        dropout_p=dropout_p,
-        is_causal=is_causal,
-        return_debug_mask=False,
-        scale=scale,
+    out, lse, cum_seq_q, cum_seq_k, max_q, max_k, philox_seed, philox_offset, debug_attn_mask = (
+        torch.ops.aten._scaled_dot_product_cudnn_attention(
+            query=query,
+            key=key,
+            value=value,
+            attn_bias=attn_mask,
+            compute_log_sumexp=return_lse,
+            dropout_p=dropout_p,
+            is_causal=is_causal,
+            return_debug_mask=False,
+            scale=scale,
+        )
     )
 
     tensors_to_save += (out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset)
@@ -982,11 +919,7 @@ def _flash_attention_backward_op(
     **kwargs,
 ):
     query, key, value, out, lse, rng_state = ctx.saved_tensors
-    grad_query, grad_key, grad_value = (
-        torch.empty_like(query),
-        torch.empty_like(key),
-        torch.empty_like(value),
-    )
+    grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value)
 
     lse_d = _wrapped_flash_attn_backward(  # noqa: F841
         grad_out,
@@ -1210,19 +1143,7 @@ def backward(
 
         grad_query, grad_key, grad_value = (x.to(grad_out.dtype) for x in (grad_query, grad_key, grad_value))
 
-        return (
-            grad_query,
-            grad_key,
-            grad_value,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
+        return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None
 
 
 class TemplatedUlyssesAttention(torch.autograd.Function):
@@ -1317,19 +1238,7 @@ def backward(
             x.flatten(0, 1).permute(1, 2, 0, 3).contiguous() for x in (grad_query, grad_key, grad_value)
         )
 
-        return (
-            grad_query,
-            grad_key,
-            grad_value,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
+        return grad_query, grad_key, grad_value, None, None, None, None, None, None, None, None
 
 
 def _templated_context_parallel_attention(
@@ -1677,12 +1586,7 @@ def _native_flex_attention(
         block_mask = attn_mask
     elif is_causal:
         block_mask = flex_attention.create_block_mask(
-            _flex_attention_causal_mask_mod,
-            batch_size,
-            num_heads,
-            seq_len_q,
-            seq_len_kv,
-            query.device,
+            _flex_attention_causal_mask_mod, batch_size, num_heads, seq_len_q, seq_len_kv, query.device
         )
     elif torch.is_tensor(attn_mask):
         if attn_mask.ndim == 2:
@@ -1702,7 +1606,6 @@ def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
 
             def score_mod(score, batch_idx, head_idx, q_idx, kv_idx):
                 return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx]
-
     else:
         raise ValueError("Attention mask must be either None, a BlockMask, or a 2D/4D tensor.")