注意力后端

注意力分发器（attention dispatcher）是一项实验性功能。如果您有任何反馈或遇到问题，请提交一个 issue。

Diffusers 通过其 attention dispatcher 提供了多种经过优化的注意力算法，这些算法在内存和计算效率方面表现更佳。该调度器充当路由器，用于管理和切换不同的注意力实现，并为与这些实现进行交互提供统一的接口。

请参阅下表，了解可用的注意力机制家族，并参阅可用后端部分以获取更完整的列表。

注意力机制	主要功能
FlashAttention	通过分块和重计算来最小化内存读写
ROCm 的 AI Tensor Engine	针对 AMD ROCm 加速器优化的 FlashAttention 实现
SageAttention	将注意力量化为 int8
PyTorch 原生	使用 scaled_dot_product_attention 的内置 PyTorch 实现
xFormers	支持多种注意力内核的内存高效注意力

本指南将向您展示如何设置和使用不同的注意力后端。

set_attention_backend

set_attention_backend() 方法会遍历模型中的所有模块，并设置要使用的相应注意力后端。注意力后端的设置会一直生效，直到调用 reset_attention_backend() 为止。

以下示例演示了如何启用来自 kernels 库的 FlashAttention-3 的 _flash_3_hub 实现，该实现允许您即时使用 Hub 上优化的计算内核，而无需任何配置。

FlashAttention-3 至少需要 Ampere 架构的 GPU。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
pipeline.transformer.set_attention_backend("_flash_3_hub")

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""
pipeline(prompt).images[0]

要恢复默认的注意力后端，请调用 reset_attention_backend()。

pipeline.transformer.reset_attention_backend()

attention_backend 上下文管理器

attention_backend 上下文管理器会在上下文内为模型临时设置一个注意力后端。在上下文之外，将使用默认的注意力机制（即 PyTorch 原生的缩放点积注意力）。如果您希望为流水线中的不同部分使用不同的后端，或者想要测试不同的后端，这一功能会非常有用。

import torch
from diffusers import QwenImagePipeline

pipeline = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)
prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

with attention_backend("_flash_3_hub"):
    image = pipeline(prompt).images[0]

大多数注意力后端支持 torch.compile，且不会导致图中断，可用于进一步加速推理。

信任远程内核

Hub 后端以及其他基于内核的功能（例如 GGUF 和 Nunchaku Lite）会通过 kernels 从 Hub 下载计算内核，并在本地执行其代码。

默认情况下，kernels 仅在其发布者属于 Hub 上的受信任内核发布者时才会加载内核。在 kernels-community 组织下发布的所有内核均受信任，因此 Diffusers 无需任何额外配置即可加载它们。_flash_3_hub、flash_hub、sage_hub 以及其他 Hub 注意力后端均解析到 kernels-community 仓库。

来自其他发布者的内核未经审核。加载此类内核会下载并运行 Diffusers 无法保证其安全性的代码，因此 Diffusers 默认将其禁用，除非您通过 DIFFUSERS_TRUST_REMOTE_KERNELS 环境变量明确选择启用。设置该变量后，Diffusers 会将 trust_remote_code=True 传递给 kernels，使其也能加载来自不受信任发布者的内核。

export DIFFUSERS_TRUST_REMOTE_KERNELS=true

仅在检查内核仓库后再启用此选项，因为它会授予已下载代码在您的机器上运行的权限。若未启用，从不受信任的发布者加载内核将引发错误。Diffusers 会自行执行此检查，因此该检查同样适用于早于 trust_remote_code 参数引入的 kernels<0.14.0。设置 DIFFUSERS_DISABLE_REMOTE_CODE=true 可全局禁用远程代码，并且其优先级高于 DIFFUSERS_TRUST_REMOTE_KERNELS。

检查

注意力分发器包含调试检查，可在问题引发故障之前捕获常见错误。

设备检查会验证查询、键和值张量是否位于同一设备上。
数据类型检查会确认张量具有匹配的数据类型，并且使用 bfloat16 或 float16。
形状检查用于验证张量维度，并防止将注意力掩码与因果标志混用。

通过设置 DIFFUSERS_ATTN_CHECKS 环境变量来启用这些检查。由于检查会为每个注意力操作增加额外开销，因此默认情况下是禁用的。

export DIFFUSERS_ATTN_CHECKS=yes

这些检查现在会在每次注意力操作之前运行。

import torch

query = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")  # or "mps", "xpu", "cpu"
key = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")
value = torch.randn(1, 10, 8, 64, dtype=torch.bfloat16, device="cuda")

try:
    with attention_backend("flash"):
        output = dispatch_attention_fn(query, key, value)
        print("✓ Flash Attention works with checks enabled")
except Exception as e:
    print(f"✗ Flash Attention failed: {e}")

您也可以直接配置注册表。

from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True

可用的后端

请参阅下表，了解可用的注意力后端及其变体的完整列表。

Expand

后端名称	家族	描述
`native`	PyTorch 原生	使用 PyTorch 的 scaled_dot_product_attention 的默认后端
`flex`	FlexAttention	PyTorch FlexAttention 实现
`_native_cudnn`	PyTorch 原生	CuDNN 优化的注意力
`_native_efficient`	PyTorch 原生	内存高效注意力
`_native_flash`	PyTorch 原生	PyTorch 的 FlashAttention
`_native_math`	PyTorch 原生	基于数学的注意力（回退方案）
`_native_npu`	PyTorch 原生	NPU 优化的注意力
`_native_xla`	PyTorch 原生	XLA 优化的注意力
`flash`	FlashAttention	FlashAttention-2
`flash_hub`	FlashAttention	来自内核的 FlashAttention-2
`flash_varlen`	FlashAttention	可变长度 FlashAttention
`flash_varlen_hub`	FlashAttention	来自内核的变长 FlashAttention
`aiter_fa2_hub`	ROCm 的 AI Tensor Engine	来自 kernels 的 AMD ROCm 版 FlashAttention-2
`aiter_sage_hub`	面向 ROCm 的 AI Tensor Engine	来自 kernels 的 AMD RDNA4 (gfx1201) 量化注意力 (INT8 QK, FP8 PV)
`flash_4_hub`	FlashAttention	FlashAttention-4
`_flash_3`	FlashAttention	FlashAttention-3
`_flash_varlen_3`	FlashAttention	可变长度 FlashAttention-3
`_flash_3_hub`	FlashAttention	来自 kernels 的 FlashAttention-3
`_flash_3_varlen_hub`	FlashAttention	来自内核的变长 FlashAttention-3
`sage`	SageAttention	量化注意力（INT8 QK）
`sage_hub`	SageAttention	来自内核的量化注意力（INT8 QK）
`sage_blackwell_hub`	SageAttention	来自 kernels 的 SageAttention3 FP4 注意力，适用于 SM120 Blackwell GPU
`sage_varlen`	SageAttention	可变长度 SageAttention
`_sage_qk_int8_pv_fp8_cuda`	SageAttention	INT8 QK + FP8 PV (CUDA)
`_sage_qk_int8_pv_fp8_cuda_sm90`	SageAttention	INT8 QK + FP8 PV (SM90)
`_sage_qk_int8_pv_fp16_cuda`	SageAttention	INT8 QK + FP16 PV (CUDA)
`_sage_qk_int8_pv_fp16_triton`	SageAttention	INT8 QK + FP16 PV (Triton)
`xformers`	xFormers	内存高效注意力

Update on GitHub