# 注意力后端

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

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

请参阅下表，了解可用的注意力机制家族，并参阅 [可用后端](#available-backends) 部分以获取更完整的列表。

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

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

## set_attention_backend

[set_attention_backend()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.set_attention_backend) 方法会遍历模型中的所有模块，并设置要使用的相应注意力后端。注意力后端的设置会一直生效，直到调用 [reset_attention_backend()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.reset_attention_backend) 为止。

以下示例演示了如何启用来自 [`kernels`](https://github.com/huggingface/kernels) 库的 FlashAttention-3 的 `_flash_3_hub` 实现，该实现允许您即时使用 Hub 上优化的计算内核，而无需任何配置。

> [!NOTE]
> FlashAttention-3 至少需要 Ampere 架构的 GPU。

```py
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()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.reset_attention_backend)。

```py
pipeline.transformer.reset_attention_backend()
```

## attention_backend 上下文管理器

[attention_backend](https://github.com/huggingface/diffusers/blob/5e181eddfe7e44c1444a2511b0d8e21d177850a0/src/diffusers/models/attention_dispatch.py#L225) 上下文管理器会在上下文内为模型临时设置一个注意力后端。在上下文之外，将使用默认的注意力机制（即 PyTorch 原生的缩放点积注意力）。如果您希望为流水线中的不同部分使用不同的后端，或者想要测试不同的后端，这一功能会非常有用。

```py
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]
```

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

## 信任远程内核

Hub 后端以及其他基于内核的功能（例如 [GGUF](../quantization/gguf) 和 [Nunchaku Lite](../quantization/nunchaku)）会通过 [`kernels`](https://github.com/huggingface/kernels) 从 Hub 下载计算内核，并在本地执行其代码。

默认情况下，`kernels` 仅在其发布者属于 Hub 上的受信任内核发布者时才会加载内核。在 [`kernels-community`](https://huggingface.co/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`，使其也能加载来自不受信任发布者的内核。

```bash
export DIFFUSERS_TRUST_REMOTE_KERNELS=true
```

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

## 检查

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

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

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

```bash
export DIFFUSERS_ATTN_CHECKS=yes
```

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

```py
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}")
```

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

```py
from diffusers.models.attention_dispatch import _AttentionBackendRegistry

_AttentionBackendRegistry._checks_enabled = True
```

## 可用的后端

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

展开

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