# Copyright 2025 The HuggingFace Team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

from __future__ import annotations

import contextlib

import functools

import inspect

import math

from dataclasses import dataclass

from enum import Enum

from typing import TYPE_CHECKING, Any, Callable

import torch

import torch.distributed as dist

import torch.nn.functional as F

if torch.distributed.is_available():

import torch.distributed._functional_collectives as funcol

from ..utils import (

get_logger,

is_aiter_available,

is_aiter_version,

is_flash_attn_3_available,

is_flash_attn_available,

is_flash_attn_version,

is_kernels_available,

is_kernels_version,

is_sageattention_available,

is_sageattention_version,

is_torch_npu_available,

is_torch_version,

is_torch_xla_available,

is_torch_xla_version,

is_xformers_available,

is_xformers_version,

)

from ..utils.constants import DIFFUSERS_ATTN_BACKEND, DIFFUSERS_ATTN_CHECKS

from ..utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph

from ._modeling_parallel import gather_size_by_comm

if TYPE_CHECKING:

from ._modeling_parallel import ParallelConfig

_REQUIRED_FLASH_VERSION = "2.6.3"

_REQUIRED_AITER_VERSION = "0.1.5"

_REQUIRED_SAGE_VERSION = "2.1.1"

_REQUIRED_FLEX_VERSION = "2.5.0"

_REQUIRED_XLA_VERSION = "2.2"

_REQUIRED_XFORMERS_VERSION = "0.0.29"

logger = get_logger(__name__) # pylint: disable=invalid-name

_CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION)

_CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available()

_CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION)

_CAN_USE_SAGE_ATTN = is_sageattention_available() and is_sageattention_version(">=", _REQUIRED_SAGE_VERSION)

_CAN_USE_FLEX_ATTN = is_torch_version(">=", _REQUIRED_FLEX_VERSION)

_CAN_USE_NPU_ATTN = is_torch_npu_available()

_CAN_USE_XLA_ATTN = is_torch_xla_available() and is_torch_xla_version(">=", _REQUIRED_XLA_VERSION)

_CAN_USE_XFORMERS_ATTN = is_xformers_available() and is_xformers_version(">=", _REQUIRED_XFORMERS_VERSION)

if _CAN_USE_FLASH_ATTN:

try:

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

except (ImportError, OSError, RuntimeError) as e:

# Handle ABI mismatch or other import failures gracefully.

# This can happen when flash_attn was compiled against a different PyTorch version.

logger.warning(f"flash_attn is installed but failed to import: {e}. Falling back to native PyTorch attention.")

_CAN_USE_FLASH_ATTN = False

flash_attn_func = None

flash_attn_varlen_func = None

_wrapped_flash_attn_backward = None

_wrapped_flash_attn_forward = None

else:

flash_attn_func = None

flash_attn_varlen_func = None

_wrapped_flash_attn_backward = None

_wrapped_flash_attn_forward = None

if _CAN_USE_FLASH_ATTN_3:

try:

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

except (ImportError, OSError, RuntimeError) as e:

logger.warning(f"flash_attn_3 failed to import: {e}. Falling back to native attention.")

_CAN_USE_FLASH_ATTN_3 = False

flash_attn_3_func = None

flash_attn_3_varlen_func = None

else:

flash_attn_3_func = None

flash_attn_3_varlen_func = None

if _CAN_USE_AITER_ATTN:

try:

from aiter import flash_attn_func as aiter_flash_attn_func

except (ImportError, OSError, RuntimeError) as e:

logger.warning(f"aiter failed to import: {e}. Falling back to native attention.")

_CAN_USE_AITER_ATTN = False

aiter_flash_attn_func = None

else:

aiter_flash_attn_func = None

if _CAN_USE_SAGE_ATTN:

try:

from sageattention import (

sageattn,

sageattn_qk_int8_pv_fp8_cuda,

sageattn_qk_int8_pv_fp8_cuda_sm90,

sageattn_qk_int8_pv_fp16_cuda,

sageattn_qk_int8_pv_fp16_triton,

sageattn_varlen,

)

except (ImportError, OSError, RuntimeError) as e:

logger.warning(f"sageattention failed to import: {e}. Falling back to native attention.")

_CAN_USE_SAGE_ATTN = False

sageattn = None

sageattn_qk_int8_pv_fp8_cuda = None

sageattn_qk_int8_pv_fp8_cuda_sm90 = None

sageattn_qk_int8_pv_fp16_cuda = None

sageattn_qk_int8_pv_fp16_triton = None

sageattn_varlen = None

else:

sageattn = None

sageattn_qk_int8_pv_fp16_cuda = None

sageattn_qk_int8_pv_fp16_triton = None

sageattn_qk_int8_pv_fp8_cuda = None

sageattn_qk_int8_pv_fp8_cuda_sm90 = None

sageattn_varlen = None

if _CAN_USE_FLEX_ATTN:

try:

# We cannot import the flex_attention function from the package directly because it is expected (from the

# pytorch documentation) that the user may compile it. If we import directly, we will not have access to the

# compiled function.

import torch.nn.attention.flex_attention as flex_attention

except (ImportError, OSError, RuntimeError) as e:

logger.warning(f"flex_attention failed to import: {e}. Falling back to native attention.")

_CAN_USE_FLEX_ATTN = False

flex_attention = None

else:

flex_attention = None

if _CAN_USE_NPU_ATTN:

try:

from torch_npu import npu_fusion_attention

except (ImportError, OSError, RuntimeError) as e:

logger.warning(f"torch_npu failed to import: {e}. Falling back to native attention.")

_CAN_USE_NPU_ATTN = False

npu_fusion_attention = None

else:

npu_fusion_attention = None

if _CAN_USE_XLA_ATTN:

try:

from torch_xla.experimental.custom_kernel import flash_attention as xla_flash_attention

except (ImportError, OSError, RuntimeError) as e:

logger.warning(f"torch_xla failed to import: {e}. Falling back to native attention.")

_CAN_USE_XLA_ATTN = False

xla_flash_attention = None

else:

xla_flash_attention = None

if _CAN_USE_XFORMERS_ATTN:

try:

import xformers.ops as xops

except (ImportError, OSError, RuntimeError) as e:

logger.warning(f"xformers failed to import: {e}. Falling back to native attention.")

_CAN_USE_XFORMERS_ATTN = False

xops = None

else:

xops = None

# Version guard for PyTorch compatibility - custom_op was added in PyTorch 2.4

if torch.__version__ >= "2.4.0":

_custom_op = torch.library.custom_op

_register_fake = torch.library.register_fake

else:

def custom_op_no_op(name, fn=None, /, *, mutates_args, device_types=None, schema=None):

def wrap(func):

return func

return wrap if fn is None else fn

def register_fake_no_op(op, fn=None, /, *, lib=None, _stacklevel=1):

def wrap(func):

return func

return wrap if fn is None else fn

_custom_op = custom_op_no_op

_register_fake = register_fake_no_op

# TODO(aryan): Add support for the following:

# - Sage Attention++

# - block sparse, radial and other attention methods

# - CP with sage attention, flex, xformers, other missing backends

# - Add support for normal and CP training with backends that don't support it yet

class AttentionBackendName(str, Enum):

# EAGER = "eager"

# `flash-attn`

FLASH = "flash"

FLASH_HUB = "flash_hub"

FLASH_VARLEN = "flash_varlen"

FLASH_VARLEN_HUB = "flash_varlen_hub"

FLASH_4_HUB = "flash_4_hub"

_FLASH_3 = "_flash_3"

_FLASH_VARLEN_3 = "_flash_varlen_3"

_FLASH_3_HUB = "_flash_3_hub"

_FLASH_3_VARLEN_HUB = "_flash_3_varlen_hub"

# `aiter`

AITER = "aiter"

# PyTorch native

FLEX = "flex"

NATIVE = "native"

_NATIVE_CUDNN = "_native_cudnn"

_NATIVE_EFFICIENT = "_native_efficient"

_NATIVE_FLASH = "_native_flash"

_NATIVE_MATH = "_native_math"

_NATIVE_NPU = "_native_npu"

_NATIVE_XLA = "_native_xla"

# `sageattention`

SAGE = "sage"

SAGE_HUB = "sage_hub"

SAGE_VARLEN = "sage_varlen"

_SAGE_QK_INT8_PV_FP8_CUDA = "_sage_qk_int8_pv_fp8_cuda"

_SAGE_QK_INT8_PV_FP8_CUDA_SM90 = "_sage_qk_int8_pv_fp8_cuda_sm90"

_SAGE_QK_INT8_PV_FP16_CUDA = "_sage_qk_int8_pv_fp16_cuda"

_SAGE_QK_INT8_PV_FP16_TRITON = "_sage_qk_int8_pv_fp16_triton"

# TODO: let's not add support for Sparge Attention now because it requires tuning per model

# We can look into supporting something "autotune"-ing in the future

# SPARGE = "sparge"

# `xformers`

XFORMERS = "xformers"

class _AttentionBackendRegistry:

_backends = {}

_constraints = {}

_supported_arg_names = {}

_supports_context_parallel = set()

_active_backend = AttentionBackendName(DIFFUSERS_ATTN_BACKEND)

_checks_enabled = DIFFUSERS_ATTN_CHECKS

@classmethod

def register(

cls,

backend: AttentionBackendName,

constraints: list[Callable] | None = None,

supports_context_parallel: bool = False,

):

logger.debug(f"Registering attention backend: {backend} with constraints: {constraints}")

def decorator(func):

cls._backends[backend] = func

cls._constraints[backend] = constraints or []

cls._supported_arg_names[backend] = set(inspect.signature(func).parameters.keys())

if supports_context_parallel:

cls._supports_context_parallel.add(backend.value)

return func

return decorator

@classmethod

def get_active_backend(cls):

return cls._active_backend, cls._backends[cls._active_backend]

@classmethod

def set_active_backend(cls, backend: str):

cls._active_backend = backend

@classmethod

def list_backends(cls):

return list(cls._backends.keys())

@classmethod

def _is_context_parallel_available(

cls,

backend: AttentionBackendName,

) -> bool:

supports_context_parallel = backend.value in cls._supports_context_parallel

return supports_context_parallel

@dataclass

class _HubKernelConfig:

"""Configuration for downloading and using a hub-based attention kernel."""

repo_id: str

function_attr: str

revision: str | None = None

version: int | None = None

kernel_fn: Callable | None = None

wrapped_forward_attr: str | None = None

wrapped_backward_attr: str | None = None

wrapped_forward_fn: Callable | None = None

wrapped_backward_fn: Callable | None = None

# Registry for hub-based attention kernels

_HUB_KERNELS_REGISTRY: dict["AttentionBackendName", _HubKernelConfig] = {

AttentionBackendName._FLASH_3_HUB: _HubKernelConfig(

repo_id="kernels-community/flash-attn3",

function_attr="flash_attn_func",

wrapped_forward_attr="flash_attn_interface._flash_attn_forward",

wrapped_backward_attr="flash_attn_interface._flash_attn_backward",

version=1,

),

AttentionBackendName._FLASH_3_VARLEN_HUB: _HubKernelConfig(

repo_id="kernels-community/flash-attn3",

function_attr="flash_attn_varlen_func",

version=1,

),

AttentionBackendName.FLASH_HUB: _HubKernelConfig(

repo_id="kernels-community/flash-attn2",

function_attr="flash_attn_func",

wrapped_forward_attr="flash_attn_interface._wrapped_flash_attn_forward",

wrapped_backward_attr="flash_attn_interface._wrapped_flash_attn_backward",

version=1,

),

AttentionBackendName.FLASH_VARLEN_HUB: _HubKernelConfig(

repo_id="kernels-community/flash-attn2",

function_attr="flash_attn_varlen_func",

version=1,

),

AttentionBackendName.SAGE_HUB: _HubKernelConfig(

repo_id="kernels-community/sage-attention",

function_attr="sageattn",

version=1,

),

AttentionBackendName.FLASH_4_HUB: _HubKernelConfig(

repo_id="kernels-staging/flash-attn4",

function_attr="flash_attn_func",

version=0,

),

}

@contextlib.contextmanager

def attention_backend(backend: str | AttentionBackendName = AttentionBackendName.NATIVE):

"""

Context manager to set the active attention backend.

"""

if backend not in _AttentionBackendRegistry._backends:

raise ValueError(f"Backend {backend} is not registered.")

backend = AttentionBackendName(backend)

_check_attention_backend_requirements(backend)

_maybe_download_kernel_for_backend(backend)

old_backend = _AttentionBackendRegistry._active_backend

_AttentionBackendRegistry.set_active_backend(backend)

try:

yield

finally:

_AttentionBackendRegistry.set_active_backend(old_backend)

def dispatch_attention_fn(

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attn_mask: torch.Tensor | None = None,

dropout_p: float = 0.0,

is_causal: bool = False,

scale: float | None = None,

enable_gqa: bool = False,

attention_kwargs: dict[str, Any] | None = None,

*,

backend: AttentionBackendName | None = None,

parallel_config: "ParallelConfig" | None = None,

) -> torch.Tensor:

attention_kwargs = attention_kwargs or {}

if backend is None:

# If no backend is specified, we either use the default backend (set via the DIFFUSERS_ATTN_BACKEND environment

# variable), or we use a custom backend based on whether user is using the `attention_backend` context manager

backend_name, backend_fn = _AttentionBackendRegistry.get_active_backend()

else:

backend_name = AttentionBackendName(backend)

backend_fn = _AttentionBackendRegistry._backends.get(backend_name)

kwargs = {

"query": query,

"key": key,

"value": value,

"attn_mask": attn_mask,

"dropout_p": dropout_p,

"is_causal": is_causal,

"scale": scale,

**attention_kwargs,

"_parallel_config": parallel_config,

}

# Equivalent to `is_torch_version(">=", "2.5.0")` — use module-level constant to avoid

# Dynamo tracing into the lru_cache-wrapped `is_torch_version` during torch.compile.

if _CAN_USE_FLEX_ATTN:

kwargs["enable_gqa"] = enable_gqa

if _AttentionBackendRegistry._checks_enabled:

removed_kwargs = set(kwargs) - set(_AttentionBackendRegistry._supported_arg_names[backend_name])

if removed_kwargs:

logger.warning(f"Removing unsupported arguments for attention backend {backend_name}: {removed_kwargs}.")

for check in _AttentionBackendRegistry._constraints.get(backend_name):

check(**kwargs)

kwargs = {k: v for k, v in kwargs.items() if k in _AttentionBackendRegistry._supported_arg_names[backend_name]}

return backend_fn(**kwargs)

# ===== Checks =====

# A list of very simple functions to catch common errors quickly when debugging.

def _check_attn_mask_or_causal(attn_mask: torch.Tensor | None, is_causal: bool, **kwargs) -> None:

if attn_mask is not None and is_causal:

raise ValueError("`is_causal` cannot be True when `attn_mask` is not None.")

def _check_device(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:

if query.device != key.device or query.device != value.device:

raise ValueError("Query, key, and value must be on the same device.")

if query.dtype != key.dtype or query.dtype != value.dtype:

raise ValueError("Query, key, and value must have the same dtype.")

def _check_device_cuda(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:

_check_device(query, key, value)

if query.device.type != "cuda":

raise ValueError("Query, key, and value must be on a CUDA device.")

def _check_device_cuda_atleast_smXY(major: int, minor: int) -> Callable:

def check_device_cuda(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:

_check_device_cuda(query, key, value)

if torch.cuda.get_device_capability(query.device) < (major, minor):

raise ValueError(

f"Query, key, and value must be on a CUDA device with compute capability >= {major}.{minor}."

)

return check_device_cuda

def _check_qkv_dtype_match(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:

if query.dtype != key.dtype:

raise ValueError("Query and key must have the same dtype.")

if query.dtype != value.dtype:

raise ValueError("Query and value must have the same dtype.")

def _check_qkv_dtype_bf16_or_fp16(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, **kwargs) -> None:

_check_qkv_dtype_match(query, key, value)

if query.dtype not in (torch.bfloat16, torch.float16):

raise ValueError("Query, key, and value must be either bfloat16 or float16.")

def _check_shape(

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attn_mask: torch.Tensor | None = None,

**kwargs,

) -> None:

# Expected shapes:

# query: (batch_size, seq_len_q, num_heads, head_dim)

# key: (batch_size, seq_len_kv, num_heads, head_dim)

# value: (batch_size, seq_len_kv, num_heads, head_dim)

# attn_mask: (seq_len_q, seq_len_kv) or (batch_size, seq_len_q, seq_len_kv)

# or (batch_size, num_heads, seq_len_q, seq_len_kv)

if query.shape[-1] != key.shape[-1]:

raise ValueError("Query and key must have the same head dimension.")

if key.shape[-3] != value.shape[-3]:

raise ValueError("Key and value must have the same sequence length.")

if attn_mask is not None and attn_mask.shape[-1] != key.shape[-3]:

raise ValueError("Attention mask must match the key's sequence length.")

# ===== Helper functions =====

def _check_attention_backend_requirements(backend: AttentionBackendName) -> None:

if backend in [AttentionBackendName.FLASH, AttentionBackendName.FLASH_VARLEN]:

if not _CAN_USE_FLASH_ATTN:

raise RuntimeError(

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]:

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."

)

elif backend in [

AttentionBackendName.FLASH_HUB,

AttentionBackendName.FLASH_VARLEN_HUB,

AttentionBackendName._FLASH_3_HUB,

AttentionBackendName._FLASH_3_VARLEN_HUB,

AttentionBackendName.SAGE_HUB,

AttentionBackendName.FLASH_4_HUB,

]:

if not is_kernels_available():

raise RuntimeError(

f"Backend '{backend.value}' is not usable because the `kernels` package isn't available. Please install it with `pip install kernels`."

)

if not is_kernels_version(">=", "0.12"):

raise RuntimeError(

f"Backend '{backend.value}' needs to be used with a `kernels` version of at least 0.12. Please update with `pip install -U kernels`."

)

if backend == AttentionBackendName.FLASH_4_HUB and not is_kernels_version(">=", "0.12.3"):

raise RuntimeError(

f"Backend '{backend.value}' needs to be used with a `kernels` version of at least 0.12.3. Please update with `pip install -U kernels`."

)

elif backend == AttentionBackendName.AITER:

if not _CAN_USE_AITER_ATTN:

raise RuntimeError(

f"Aiter Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `aiter>={_REQUIRED_AITER_VERSION}`."

)

elif backend in [

AttentionBackendName.SAGE,

AttentionBackendName.SAGE_VARLEN,

AttentionBackendName._SAGE_QK_INT8_PV_FP8_CUDA,

AttentionBackendName._SAGE_QK_INT8_PV_FP8_CUDA_SM90,

AttentionBackendName._SAGE_QK_INT8_PV_FP16_CUDA,

AttentionBackendName._SAGE_QK_INT8_PV_FP16_TRITON,

]:

if not _CAN_USE_SAGE_ATTN:

raise RuntimeError(

f"Sage Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `sageattention>={_REQUIRED_SAGE_VERSION}`."

)

elif backend == AttentionBackendName.FLEX:

if not _CAN_USE_FLEX_ATTN:

raise RuntimeError(

f"Flex Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `torch>=2.5.0`."

)

elif backend == AttentionBackendName._NATIVE_NPU:

if not _CAN_USE_NPU_ATTN:

raise RuntimeError(

f"NPU Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `torch_npu`."

)

elif backend == AttentionBackendName._NATIVE_XLA:

if not _CAN_USE_XLA_ATTN:

raise RuntimeError(

f"XLA Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `torch_xla>={_REQUIRED_XLA_VERSION}`."

)

elif backend == AttentionBackendName.XFORMERS:

if not _CAN_USE_XFORMERS_ATTN:

raise RuntimeError(

f"Xformers Attention backend '{backend.value}' is not usable because of missing package or the version is too old. Please install `xformers>={_REQUIRED_XFORMERS_VERSION}`."

)

@lru_cache_unless_export(maxsize=128)

def _prepare_for_flash_attn_or_sage_varlen_without_mask(

batch_size: int,

seq_len_q: int,

seq_len_kv: int,

device: torch.device | None = None,

):

seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device)

seqlens_k = torch.full((batch_size,), seq_len_kv, dtype=torch.int32, device=device)

cu_seqlens_q = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)

cu_seqlens_k = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)

cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0)

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)

def _prepare_for_flash_attn_or_sage_varlen_with_mask(

batch_size: int,

seq_len_q: int,

attn_mask: torch.Tensor,

device: torch.device | None = None,

):

seqlens_q = torch.full((batch_size,), seq_len_q, dtype=torch.int32, device=device)

seqlens_k = attn_mask.sum(dim=1, dtype=torch.int32)

cu_seqlens_q = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)

cu_seqlens_k = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)

cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0)

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)

def _prepare_for_flash_attn_or_sage_varlen(

batch_size: int,

seq_len_q: int,

seq_len_kv: int,

attn_mask: torch.Tensor | None = None,

device: torch.device | None = None,

) -> None:

if attn_mask is None:

return _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, device)

return _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, device)

def _normalize_attn_mask(attn_mask: torch.Tensor, batch_size: int, seq_len_k: int) -> torch.Tensor:

"""

Normalize an attention mask to shape [batch_size, seq_len_k] (bool) suitable for inferring seqlens_[q|k] in

FlashAttention/Sage varlen.

Supports 1D to 4D shapes and common broadcasting patterns.

"""

if attn_mask.dtype != torch.bool:

raise ValueError(f"Attention mask must be of type bool, got {attn_mask.dtype}.")

if attn_mask.ndim == 1:

# [seq_len_k] -> broadcast across batch

attn_mask = attn_mask.unsqueeze(0).expand(batch_size, seq_len_k)

elif attn_mask.ndim == 2:

# [batch_size, seq_len_k]. Maybe broadcast across batch

if attn_mask.size(0) not in [1, batch_size]:

raise ValueError(

f"attn_mask.shape[0] ({attn_mask.shape[0]}) must be 1 or {batch_size} for 2D attention mask."

)

attn_mask = attn_mask.expand(batch_size, seq_len_k)

elif attn_mask.ndim == 3:

# [batch_size, seq_len_q, seq_len_k] -> reduce over query dimension

# We do this reduction because we know that arbitrary QK masks is not supported in Flash/Sage varlen.

if attn_mask.size(0) not in [1, batch_size]:

raise ValueError(

f"attn_mask.shape[0] ({attn_mask.shape[0]}) must be 1 or {batch_size} for 3D attention mask."

)

attn_mask = attn_mask.any(dim=1)

attn_mask = attn_mask.expand(batch_size, seq_len_k)

elif attn_mask.ndim == 4:

# [batch_size, num_heads, seq_len_q, seq_len_k] or broadcastable versions

if attn_mask.size(0) not in [1, batch_size]:

raise ValueError(

f"attn_mask.shape[0] ({attn_mask.shape[0]}) must be 1 or {batch_size} for 4D attention mask."

)

attn_mask = attn_mask.expand(batch_size, -1, -1, seq_len_k) # [B, H, Q, K]

attn_mask = attn_mask.any(dim=(1, 2)) # [B, K]

else:

raise ValueError(f"Unsupported attention mask shape: {attn_mask.shape}")

if attn_mask.shape != (batch_size, seq_len_k):

raise ValueError(

f"Normalized attention mask shape mismatch: got {attn_mask.shape}, expected ({batch_size}, {seq_len_k})"

)

return attn_mask

def _flex_attention_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):

return q_idx >= kv_idx

# ===== Helpers for downloading kernels =====

def _resolve_kernel_attr(module, attr_path: str):

target = module

for attr in attr_path.split("."):

if not hasattr(target, attr):

raise AttributeError(f"Kernel module '{module.__name__}' does not define attribute path '{attr_path}'.")

target = getattr(target, attr)

return target

def _maybe_download_kernel_for_backend(backend: AttentionBackendName) -> None:

if backend not in _HUB_KERNELS_REGISTRY:

return

config = _HUB_KERNELS_REGISTRY[backend]

needs_kernel = config.kernel_fn is None

needs_wrapped_forward = config.wrapped_forward_attr is not None and config.wrapped_forward_fn is None

needs_wrapped_backward = config.wrapped_backward_attr is not None and config.wrapped_backward_fn is None

if not (needs_kernel or needs_wrapped_forward or needs_wrapped_backward):

return

try:

from kernels import get_kernel

kernel_module = get_kernel(config.repo_id, revision=config.revision, version=config.version)

if needs_kernel:

config.kernel_fn = _resolve_kernel_attr(kernel_module, config.function_attr)

if needs_wrapped_forward:

config.wrapped_forward_fn = _resolve_kernel_attr(kernel_module, config.wrapped_forward_attr)

if needs_wrapped_backward:

config.wrapped_backward_fn = _resolve_kernel_attr(kernel_module, config.wrapped_backward_attr)

except Exception as e:

logger.error(f"An error occurred while fetching kernel '{config.repo_id}' from the Hub: {e}")

raise

# ===== torch op registrations =====

# Registrations are required for fullgraph tracing compatibility

# TODO: this is only required because the beta release FA3 does not have it. There is a PR adding

# this but it was never merged: https://github.com/Dao-AILab/flash-attention/pull/1590

@_custom_op("_diffusers_flash_attn_3::_flash_attn_forward", mutates_args=(), device_types="cuda")

def _wrapped_flash_attn_3(

q: torch.Tensor,

k: torch.Tensor,

v: torch.Tensor,

softmax_scale: float | None = None,

causal: bool = False,

qv: torch.Tensor | None = None,

q_descale: torch.Tensor | None = None,

k_descale: torch.Tensor | None = None,

v_descale: torch.Tensor | None = None,

attention_chunk: int = 0,

softcap: float = 0.0,

num_splits: int = 1,

pack_gqa: bool | None = None,

deterministic: bool = False,

sm_margin: int = 0,

) -> tuple[torch.Tensor, torch.Tensor]:

# Hardcoded for now because pytorch does not support tuple/int type hints

window_size = (-1, -1)

result = flash_attn_3_func(

q=q,

k=k,

v=v,

softmax_scale=softmax_scale,

causal=causal,

qv=qv,

q_descale=q_descale,

k_descale=k_descale,

v_descale=v_descale,

window_size=window_size,

attention_chunk=attention_chunk,

softcap=softcap,

num_splits=num_splits,

pack_gqa=pack_gqa,

deterministic=deterministic,

sm_margin=sm_margin,

return_attn_probs=True,

)

out, lse, *_ = result

lse = lse.permute(0, 2, 1)

return out, lse

@_register_fake("_diffusers_flash_attn_3::_flash_attn_forward")

def _(

q: torch.Tensor,

k: torch.Tensor,

v: torch.Tensor,

softmax_scale: float | None = None,

causal: bool = False,

qv: torch.Tensor | None = None,

q_descale: torch.Tensor | None = None,

k_descale: torch.Tensor | None = None,

v_descale: torch.Tensor | None = None,

attention_chunk: int = 0,

softcap: float = 0.0,

num_splits: int = 1,

pack_gqa: bool | None = None,

deterministic: bool = False,

sm_margin: int = 0,

) -> tuple[torch.Tensor, torch.Tensor]:

window_size = (-1, -1) # noqa: F841

# A lot of the parameters here are not yet used in any way within diffusers.

# We can safely ignore for now and keep the fake op shape propagation simple.

batch_size, seq_len, num_heads, head_dim = q.shape

lse_shape = (batch_size, seq_len, num_heads)

return torch.empty_like(q), q.new_empty(lse_shape)

# ===== Helper functions to use attention backends with templated CP autograd functions =====

def _native_attention_forward_op(

ctx: torch.autograd.function.FunctionCtx,

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attn_mask: torch.Tensor | None = None,

dropout_p: float = 0.0,

is_causal: bool = False,

scale: float | None = None,

enable_gqa: bool = False,

return_lse: bool = False,

_save_ctx: bool = True,

_parallel_config: "ParallelConfig" | None = None,

):

# Native attention does not return_lse

if return_lse:

raise ValueError("Native attention does not support return_lse=True")

# used for backward pass

if _save_ctx:

ctx.save_for_backward(query, key, value)

ctx.attn_mask = attn_mask

ctx.dropout_p = dropout_p

ctx.is_causal = is_causal

ctx.scale = scale

ctx.enable_gqa = enable_gqa

query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))

out = torch.nn.functional.scaled_dot_product_attention(

query=query,

key=key,

value=value,

attn_mask=attn_mask,

dropout_p=dropout_p,

is_causal=is_causal,

scale=scale,

enable_gqa=enable_gqa,

)

out = out.permute(0, 2, 1, 3)

return out

def _native_attention_backward_op(

ctx: torch.autograd.function.FunctionCtx,

grad_out: torch.Tensor,

*args,

**kwargs,

):

query, key, value = ctx.saved_tensors

query.requires_grad_(True)

key.requires_grad_(True)

value.requires_grad_(True)

with torch.enable_grad():

query_t, key_t, value_t = (x.permute(0, 2, 1, 3) for x in (query, key, value))

out = torch.nn.functional.scaled_dot_product_attention(

query=query_t,

key=key_t,

value=value_t,

attn_mask=ctx.attn_mask,

dropout_p=ctx.dropout_p,

is_causal=ctx.is_causal,

scale=ctx.scale,

enable_gqa=ctx.enable_gqa,

)

out = 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, retain_graph=False

)

grad_query = grad_query_t.permute(0, 2, 1, 3)

grad_key = grad_key_t.permute(0, 2, 1, 3)

grad_value = grad_value_t.permute(0, 2, 1, 3)

return grad_query, grad_key, grad_value

# https://github.com/pytorch/pytorch/blob/8904ba638726f8c9a5aff5977c4aa76c9d2edfa6/aten/src/ATen/native/native_functions.yaml#L14958

# forward declaration:

# aten::_scaled_dot_product_cudnn_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, bool compute_log_sumexp, float dropout_p=0., bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask)

def _cudnn_attention_forward_op(

ctx: torch.autograd.function.FunctionCtx,

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attn_mask: torch.Tensor | None = None,

dropout_p: float = 0.0,

is_causal: bool = False,

scale: float | None = None,

enable_gqa: bool = False,

return_lse: bool = False,

_save_ctx: bool = True,

_parallel_config: "ParallelConfig" | None = None,

):

if enable_gqa:

raise ValueError("`enable_gqa` is not yet supported for cuDNN attention.")

tensors_to_save = ()

# Contiguous is a must here! Calling cuDNN backend with aten ops produces incorrect results

# if the input tensors are not contiguous.

query = query.transpose(1, 2).contiguous()

key = key.transpose(1, 2).contiguous()

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,

)

)

tensors_to_save += (out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset)

if _save_ctx:

ctx.save_for_backward(*tensors_to_save)

ctx.dropout_p = dropout_p

ctx.is_causal = is_causal

ctx.scale = scale

ctx.attn_mask = attn_mask

ctx.max_q = max_q

ctx.max_k = max_k

out = out.transpose(1, 2).contiguous()

if lse is not None:

lse = lse.transpose(1, 2).contiguous()

return (out, lse) if return_lse else out

# backward declaration:

# aten::_scaled_dot_product_cudnn_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor philox_seed, Tensor philox_offset, Tensor attn_bias, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, *, float? scale=None) -> (Tensor, Tensor, Tensor)

def _cudnn_attention_backward_op(

ctx: torch.autograd.function.FunctionCtx,

grad_out: torch.Tensor,

*args,

**kwargs,

):

query, key, value, out, lse, cum_seq_q, cum_seq_k, philox_seed, philox_offset = ctx.saved_tensors

grad_out = grad_out.transpose(1, 2).contiguous()

key = key.transpose(1, 2).contiguous()

value = value.transpose(1, 2).contiguous()

# Cannot pass first 5 arguments as kwargs because: https://github.com/pytorch/pytorch/blob/d26ca5de058dbcf56ac52bb43e84dd98df2ace97/torch/_dynamo/variables/torch.py#L1341

grad_query, grad_key, grad_value = torch.ops.aten._scaled_dot_product_cudnn_attention_backward(

grad_out,

query,

key,

value,

out,

logsumexp=lse,

philox_seed=philox_seed,

philox_offset=philox_offset,

attn_bias=ctx.attn_mask,

cum_seq_q=cum_seq_q,

cum_seq_k=cum_seq_k,

max_q=ctx.max_q,

max_k=ctx.max_k,

dropout_p=ctx.dropout_p,

is_causal=ctx.is_causal,

scale=ctx.scale,

)

grad_query, grad_key, grad_value = (x.transpose(1, 2).contiguous() for x in (grad_query, grad_key, grad_value))

return grad_query, grad_key, grad_value

# https://github.com/pytorch/pytorch/blob/e33fa0ece36a93dbc8ff19b0251b8d99f8ae8668/aten/src/ATen/native/native_functions.yaml#L15135

# forward declaration:

# aten::_scaled_dot_product_flash_attention(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor rng_state, Tensor unused, Tensor debug_attn_mask)

def _native_flash_attention_forward_op(

ctx: torch.autograd.function.FunctionCtx,

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attn_mask: torch.Tensor | None = None,

dropout_p: float = 0.0,

is_causal: bool = False,

scale: float | None = None,