Diffusers Attention Backends 实战指南用注意力调度器切换 FlashAttention、SageAttention 等优化后端【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本指南围绕 Diffusers 的Attention Dispatcher注意力调度器展开它是 diffusers 中用于管理、路由和切换多种注意力Attention实现的核心机制。你将学会通过set_attention_backend、attention_backend上下文管理器为模型配置 FlashAttention、SageAttention、PyTorch native、xFormers 等后端掌握远程内核信任机制与调试检查开关并对照源码理解调度器的工作方式。[!NOTE] 注意力调度器attention dispatcher属于实验性功能API 未来可能调整。如果遇到问题或有反馈请在 diffusers 仓库中提交 issue。什么是 Attention DispatcherDiffusers 通过attention dispatcher提供多种在显存与计算效率上更优的注意力算法。调度器充当一个路由器router负责管理和切换不同的注意力实现同时对外提供统一接口让上层模型代码无需关心底层到底跑的是哪个 kernel。从源码结构看整个机制集中在 attention_dispatch.pyAttentionBackendName枚举第 215-255 行定义了全部受支持后端的规范化名称_AttentionBackendRegistry注册表第 258-304 行维护了后端函数、约束条件、支持的参数名、上下文并行支持情况以及当前激活的后端dispatch_attention_fn第 391-440 行是统一的分发入口任何注意力调用最终都汇聚到这里根据当前激活的后端把query/key/value等参数转发给对应的实现。当前主要包含以下几类注意力家族attention family主要特性FlashAttention通过 tiling分块与 recomputation重计算最小化显存读写AI Tensor Engine for ROCm专为 AMD ROCm 加速器优化的 FlashAttention 实现SageAttention将注意力量化为 int8 以加速PyTorch nativePyTorch 内置实现基于 scaled_dot_product_attentionxFormers支持多种注意力 kernel 的高显存效率注意力用 set_attention_backend 设置持久化后端[~ModelMixin.set_attention_backend] 方法会遍历模型中的所有子模块为其设置合适的注意力后端。该设置会持续生效直到调用 [~ModelMixin.reset_attention_backend] 恢复默认。下面的示例为 Qwen-Image 的 transformer 启用来自kernels库的 FlashAttention-3 实现_flash_3_hub。kernels允许直接从 Hub 获取开箱即用的优化计算内核无需任何额外编译或安装。[!NOTE] FlashAttention-3 至少需要 Ampere 架构及以上的 GPU。import torch from diffusers import QwenImagePipeline pipeline QwenImagePipeline.from_pretrained( Qwen/Qwen-Image, dtypetorch.bfloat16, device_mapcuda # 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()源码视角set_attention_backend 内部做了什么在 modeling_utils.py 中set_attention_backend第 599-659 行的实现依次完成以下工作遍历模型模块只处理Attention、MochiAttention以及AttentionModuleMixin三类注意力模块检查上下文并行兼容性如果模型已启用上下文并行context parallelism且所选后端不支持 CP会直接抛出ValueError并提示可用的兼容后端列表由注册表_supports_context_parallel集合决定校验后端名称backend 名称会被统一转小写必须在AttentionBackendName枚举范围内否则报错检查依赖与版本调用_check_attention_backend_requirements验证对应包是否安装且版本达标详见下文依赖与版本门槛按需下载远程内核对于_flash_3_hub、flash_hub等 Hub 类后端调用_maybe_download_kernel_for_backend从 Hub 拉取 kernel写入 processor将选中后端写入每个注意力模块 processor 的_attention_backend属性并同步更新注册表全局激活后端。reset_attention_backend第 661-679 行则把所有 processor 的_attention_backend置为None此后模型将回到环境变量指定的默认值未设置时即 PyTorch 原生 SDPA。用 attention_backend 上下文管理器做临时切换attention_backend 上下文管理器可以在上下文内临时为模型设置一个注意力后端退出上下文后恢复默认注意力即 PyTorch 原生 scaled dot product attention。这在需要为流水线的不同部分使用不同后端或想快速对比多个后端效果时非常有用。import torch from diffusers import QwenImagePipeline from diffusers.models import attention_backend pipeline QwenImagePipeline.from_pretrained( Qwen/Qwen-Image, dtypetorch.bfloat16, device_mapcuda # 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]其底层实现attention_dispatch.py 第 370-388 行很直观进入上下文时校验后端已注册、检查依赖版本、必要时下载 kernel然后保存旧后端并激活新后端finally块中恢复旧后端保证即使中途抛异常也不会污染全局状态。[!TIP] 大多数注意力后端都支持torch.compile且不会产生 graph break可在此基础上进一步加速推理。源码中dispatch_attention_fn大量使用模块级常量如_CUDNN_LSE_HAS_TRAILING_DIM以避免 Dynamo 追踪时产生 graph break见 attention_dispatch.py 第 73-76 行就是为此专门设计的。依赖与版本门槛每个后端都有前置依赖与最低版本要求这些约束集中在_check_attention_backend_requirementsattention_dispatch.py 第 519-590 行中切换后端时由set_attention_backend与attention_backend统一执行校验不满足会抛出RuntimeError并给出安装提示后端族前置依赖最低版本flash/flash_varlenflash-attn 2.6.3_flash_3/_flash_varlen_3FlashAttention-3需从源码构建 beta 版—flash_hub/flash_varlen_hub/_flash_3_hub/_flash_3_varlen_hub/sage_hub/flash_4_hub/aiter_fa2_hubkernels包 0.12flash_4_hub需 0.12.3sage/sage_varlen及_sage_*变体sageattention 2.1.1flextorchFlexAttention 2.5.0_native_nputorch_npu—_native_xlatorch_xla 2.2xformersxformers 0.0.29从源码结构看attention_dispatch.py会在模块导入时通过is_flash_attn_available()、is_sageattention_available()等检查第 66-72 行探测各依赖是否可用即便安装过但导入失败如 ABI 不匹配也会记录 warning 并回退到原生注意力。信任远程内核Trusting remote kernelsHub 后端以及其他依赖 kernel 的特性如 GGUF 与 Nunchaku Lite会通过kernels库从 Hub 下载计算内核并在本地执行其代码。默认情况下kernels只加载发布者是 Hub 上受信任内核发布者的内核。发布在kernels-community组织下的内核被视为受信任Diffusers 无需额外配置即可加载。_flash_3_hub、flash_hub、sage_hub以及其他 Hub 注意力后端解析到的都是kernels-community下的仓库——这一点可以在_HUB_KERNELS_REGISTRYattention_dispatch.py 第 322-367 行中直接看到例如_flash_3_hub→kernels-community/flash-attn3flash_hub/flash_varlen_hub→kernels-community/flash-attn2sage_hub→kernels-community/sage-attentionflash_4_hub→kernels-community/flash-attn4aiter_fa2_hub→kernels-community/aiter-flash-attn-ck来自其他发布者的 kernel 未经审查。加载它意味着下载并在你的机器上运行 Diffusers 无法担保的代码因此 Diffusers 默认保持禁用除非你显式通过DIFFUSERS_TRUST_REMOTE_KERNELS环境变量选择开启。设置后Diffusers 会向kernels传递trust_remote_codeTrue从而允许其加载来自非受信任发布者的内核export DIFFUSERS_TRUST_REMOTE_KERNELStrue⚠️安全提醒仅在仔细检查过内核仓库后再开启此选项因为它赋予下载的代码在你的机器上运行的权限。不设置时从非受信任发布者加载内核会直接报错。该检查由 Diffusers 自身完成因此同样适用于早于trust_remote_code参数的kernels0.14.0版本。此外设置DIFFUSERS_DISABLE_REMOTE_CODEtrue会全局禁用远程代码且其优先级高于DIFFUSERS_TRUST_REMOTE_KERNELS。这一点在 constants.py 第 49-55 行 中有直接体现DIFFUSERS_TRUST_REMOTE_KERNELS的计算条件是环境变量为真且未设置DIFFUSERS_DISABLE_REMOTE_CODE。调试检查Checks注意力调度器内置了若干调试检查用于在问题酿成大错前捕获常见错误。默认关闭可通过DIFFUSERS_ATTN_CHECKS环境变量开启export DIFFUSERS_ATTN_CHECKSyes开启后每次注意力运算前都会执行检查。检查主要分三类实现见 attention_dispatch.py 第 443-514 行设备检查device checks验证 query、key、value 张量位于同一设备_check_device部分后端还会额外要求位于 CUDA 设备或满足特定算力_check_device_cuda_atleast_smXY数据类型检查dtype checks确认 Q/K/V 的 dtype 一致且为 bfloat16 或 float16_check_qkv_dtype_bf16_or_fp16个别后端如 FA3进一步要求 bfloat16形状检查shape checks校验张量维度例如 query 与 key 的 head 维度一致、key 与 value 的序列长度一致、attention mask 与 key 的序列长度匹配并阻止同时混用 attention mask 与 causal 标志_check_attn_mask_or_causal当attn_mask不为空且is_causalTrue时直接抛错。开启检查后可用下面的片段直接验证某个后端在当前环境下是否可用import torch query torch.randn(1, 10, 8, 64, dtypetorch.bfloat16, devicecuda) # or mps, xpu, cpu key torch.randn(1, 10, 8, 64, dtypetorch.bfloat16, devicecuda) value torch.randn(1, 10, 8, 64, dtypetorch.bfloat16, devicecuda) 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})[!NOTE] 检查会给每一次注意力操作带来额外开销所以默认是关闭的建议在调试阶段开启正式跑推理时保持关闭以获得最佳性能。也可以直接配置注册表来开启检查from diffusers.models.attention_dispatch import _AttentionBackendRegistry _AttentionBackendRegistry._checks_enabled True需要注意dispatch_attention_fn内部开启检查时第 431-440 行还会自动剔除当前后端不支持的额外参数并给出 warning例如启用 GQA 时只有支持enable_gqa的后端才会收到该参数。Available backends 完整列表下表列出了当前版本可用的全部注意力后端及其变体与AttentionBackendName枚举一一对应见 attention_dispatch.py 第 215-255 行后端名称家族说明nativePyTorch native默认后端使用 PyTorch 的 scaled_dot_product_attentionflexFlexAttentionPyTorch FlexAttention 实现_native_cudnnPyTorch nativeCuDNN 优化注意力_native_efficientPyTorch native显存高效注意力_native_flashPyTorch nativePyTorch 的 FlashAttention_native_mathPyTorch native基于数学计算的注意力回退_native_npuPyTorch nativeNPU 优化注意力_native_xlaPyTorch nativeXLA 优化注意力flashFlashAttentionFlashAttention-2flash_hubFlashAttention来自 kernels 的 FlashAttention-2flash_varlenFlashAttention变长序列 FlashAttentionflash_varlen_hubFlashAttention来自 kernels 的变长 FlashAttentionaiter_fa2_hubAI Tensor Engine for ROCm来自 kernels 的 AMD ROCm FlashAttention-2flash_4_hubFlashAttentionFlashAttention-4_flash_3FlashAttentionFlashAttention-3_flash_varlen_3FlashAttention变长 FlashAttention-3_flash_3_hubFlashAttention来自 kernels 的 FlashAttention-3_flash_3_varlen_hubFlashAttention来自 kernels 的变长 FlashAttention-3sageSageAttention量化注意力INT8 QKsage_hubSageAttention来自 kernels 的量化注意力INT8 QKsage_varlenSageAttention变长 SageAttention_sage_qk_int8_pv_fp8_cudaSageAttentionINT8 QK FP8 PVCUDA_sage_qk_int8_pv_fp8_cuda_sm90SageAttentionINT8 QK FP8 PVSM90 架构_sage_qk_int8_pv_fp16_cudaSageAttentionINT8 QK FP16 PVCUDA_sage_qk_int8_pv_fp16_tritonSageAttentionINT8 QK FP16 PVTritonxformersxFormers显存高效注意力dispatch_attention_fn会以注册表记录的_supported_arg_names由函数签名自动推导为基准做参数过滤因此各后端只需实现自己支持的参数子集即可无缝接入见 attention_dispatch.py 第 278 行 与 第 438 行。如何选择与验证后端结合文档与源码给出如下实践建议追求零配置开箱即用优先选择native默认或 Hub 类后端_flash_3_hub、flash_hub、sage_hub等它们不需要本地编译 CUDA 扩展只需pip install kernels即可自动拉取预编译内核显存受限的大模型推理优先尝试 FlashAttention 系列flash/flash_varlen或 xFormersxformers两者均以显存效率著称SageAttention 的 int8 量化路径sage/_sage_*适合进一步降低显存与带宽开销AMD ROCm 平台使用aiter_fa2_hub获得针对 ROCm 优化的 FlashAttention-2对比与回归用attention_backend上下文管理器在同一个推理脚本里逐个切换后端对比速度与显存占用开启DIFFUSERS_ATTN_CHECKS便于在切换时第一时间捕获设备、dtype、形状等基础错误验证正确性仓库在 tests/models/test_attention_dispatch.py 中提供了注意力后端的相关测试如上下文并行下的test_attn_backend_backward_parity可参考其思路做不同后端输出的数值一致性对比。另外注意set_attention_backend与reset_attention_backend是ModelMixin的方法适用于任意继承它的模型如pipeline.transformer而attention_backend上下文管理器与dispatch_attention_fn均可从diffusers.models导入见 models/init.py 第 174 行 与 diffusers/init.py 第 369 行。若希望全局指定默认后端可在启动时设置DIFFUSERS_ATTN_BACKEND环境变量默认值native见 constants.py 第 45 行注册表的_active_backend会直接读取该变量作为初始值attention_dispatch.py 第 263 行。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Diffusers Attention Backends 实战指南：用注意力调度器切换 FlashAttention、SageAttention 等优化后端
Diffusers Attention Backends 实战指南用注意力调度器切换 FlashAttention、SageAttention 等优化后端【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers本指南围绕 Diffusers 的Attention Dispatcher注意力调度器展开它是 diffusers 中用于管理、路由和切换多种注意力Attention实现的核心机制。你将学会通过set_attention_backend、attention_backend上下文管理器为模型配置 FlashAttention、SageAttention、PyTorch native、xFormers 等后端掌握远程内核信任机制与调试检查开关并对照源码理解调度器的工作方式。[!NOTE] 注意力调度器attention dispatcher属于实验性功能API 未来可能调整。如果遇到问题或有反馈请在 diffusers 仓库中提交 issue。什么是 Attention DispatcherDiffusers 通过attention dispatcher提供多种在显存与计算效率上更优的注意力算法。调度器充当一个路由器router负责管理和切换不同的注意力实现同时对外提供统一接口让上层模型代码无需关心底层到底跑的是哪个 kernel。从源码结构看整个机制集中在 attention_dispatch.pyAttentionBackendName枚举第 215-255 行定义了全部受支持后端的规范化名称_AttentionBackendRegistry注册表第 258-304 行维护了后端函数、约束条件、支持的参数名、上下文并行支持情况以及当前激活的后端dispatch_attention_fn第 391-440 行是统一的分发入口任何注意力调用最终都汇聚到这里根据当前激活的后端把query/key/value等参数转发给对应的实现。当前主要包含以下几类注意力家族attention family主要特性FlashAttention通过 tiling分块与 recomputation重计算最小化显存读写AI Tensor Engine for ROCm专为 AMD ROCm 加速器优化的 FlashAttention 实现SageAttention将注意力量化为 int8 以加速PyTorch nativePyTorch 内置实现基于 scaled_dot_product_attentionxFormers支持多种注意力 kernel 的高显存效率注意力用 set_attention_backend 设置持久化后端[~ModelMixin.set_attention_backend] 方法会遍历模型中的所有子模块为其设置合适的注意力后端。该设置会持续生效直到调用 [~ModelMixin.reset_attention_backend] 恢复默认。下面的示例为 Qwen-Image 的 transformer 启用来自kernels库的 FlashAttention-3 实现_flash_3_hub。kernels允许直接从 Hub 获取开箱即用的优化计算内核无需任何额外编译或安装。[!NOTE] FlashAttention-3 至少需要 Ampere 架构及以上的 GPU。import torch from diffusers import QwenImagePipeline pipeline QwenImagePipeline.from_pretrained( Qwen/Qwen-Image, dtypetorch.bfloat16, device_mapcuda # 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()源码视角set_attention_backend 内部做了什么在 modeling_utils.py 中set_attention_backend第 599-659 行的实现依次完成以下工作遍历模型模块只处理Attention、MochiAttention以及AttentionModuleMixin三类注意力模块检查上下文并行兼容性如果模型已启用上下文并行context parallelism且所选后端不支持 CP会直接抛出ValueError并提示可用的兼容后端列表由注册表_supports_context_parallel集合决定校验后端名称backend 名称会被统一转小写必须在AttentionBackendName枚举范围内否则报错检查依赖与版本调用_check_attention_backend_requirements验证对应包是否安装且版本达标详见下文依赖与版本门槛按需下载远程内核对于_flash_3_hub、flash_hub等 Hub 类后端调用_maybe_download_kernel_for_backend从 Hub 拉取 kernel写入 processor将选中后端写入每个注意力模块 processor 的_attention_backend属性并同步更新注册表全局激活后端。reset_attention_backend第 661-679 行则把所有 processor 的_attention_backend置为None此后模型将回到环境变量指定的默认值未设置时即 PyTorch 原生 SDPA。用 attention_backend 上下文管理器做临时切换attention_backend 上下文管理器可以在上下文内临时为模型设置一个注意力后端退出上下文后恢复默认注意力即 PyTorch 原生 scaled dot product attention。这在需要为流水线的不同部分使用不同后端或想快速对比多个后端效果时非常有用。import torch from diffusers import QwenImagePipeline from diffusers.models import attention_backend pipeline QwenImagePipeline.from_pretrained( Qwen/Qwen-Image, dtypetorch.bfloat16, device_mapcuda # 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]其底层实现attention_dispatch.py 第 370-388 行很直观进入上下文时校验后端已注册、检查依赖版本、必要时下载 kernel然后保存旧后端并激活新后端finally块中恢复旧后端保证即使中途抛异常也不会污染全局状态。[!TIP] 大多数注意力后端都支持torch.compile且不会产生 graph break可在此基础上进一步加速推理。源码中dispatch_attention_fn大量使用模块级常量如_CUDNN_LSE_HAS_TRAILING_DIM以避免 Dynamo 追踪时产生 graph break见 attention_dispatch.py 第 73-76 行就是为此专门设计的。依赖与版本门槛每个后端都有前置依赖与最低版本要求这些约束集中在_check_attention_backend_requirementsattention_dispatch.py 第 519-590 行中切换后端时由set_attention_backend与attention_backend统一执行校验不满足会抛出RuntimeError并给出安装提示后端族前置依赖最低版本flash/flash_varlenflash-attn 2.6.3_flash_3/_flash_varlen_3FlashAttention-3需从源码构建 beta 版—flash_hub/flash_varlen_hub/_flash_3_hub/_flash_3_varlen_hub/sage_hub/flash_4_hub/aiter_fa2_hubkernels包 0.12flash_4_hub需 0.12.3sage/sage_varlen及_sage_*变体sageattention 2.1.1flextorchFlexAttention 2.5.0_native_nputorch_npu—_native_xlatorch_xla 2.2xformersxformers 0.0.29从源码结构看attention_dispatch.py会在模块导入时通过is_flash_attn_available()、is_sageattention_available()等检查第 66-72 行探测各依赖是否可用即便安装过但导入失败如 ABI 不匹配也会记录 warning 并回退到原生注意力。信任远程内核Trusting remote kernelsHub 后端以及其他依赖 kernel 的特性如 GGUF 与 Nunchaku Lite会通过kernels库从 Hub 下载计算内核并在本地执行其代码。默认情况下kernels只加载发布者是 Hub 上受信任内核发布者的内核。发布在kernels-community组织下的内核被视为受信任Diffusers 无需额外配置即可加载。_flash_3_hub、flash_hub、sage_hub以及其他 Hub 注意力后端解析到的都是kernels-community下的仓库——这一点可以在_HUB_KERNELS_REGISTRYattention_dispatch.py 第 322-367 行中直接看到例如_flash_3_hub→kernels-community/flash-attn3flash_hub/flash_varlen_hub→kernels-community/flash-attn2sage_hub→kernels-community/sage-attentionflash_4_hub→kernels-community/flash-attn4aiter_fa2_hub→kernels-community/aiter-flash-attn-ck来自其他发布者的 kernel 未经审查。加载它意味着下载并在你的机器上运行 Diffusers 无法担保的代码因此 Diffusers 默认保持禁用除非你显式通过DIFFUSERS_TRUST_REMOTE_KERNELS环境变量选择开启。设置后Diffusers 会向kernels传递trust_remote_codeTrue从而允许其加载来自非受信任发布者的内核export DIFFUSERS_TRUST_REMOTE_KERNELStrue⚠️安全提醒仅在仔细检查过内核仓库后再开启此选项因为它赋予下载的代码在你的机器上运行的权限。不设置时从非受信任发布者加载内核会直接报错。该检查由 Diffusers 自身完成因此同样适用于早于trust_remote_code参数的kernels0.14.0版本。此外设置DIFFUSERS_DISABLE_REMOTE_CODEtrue会全局禁用远程代码且其优先级高于DIFFUSERS_TRUST_REMOTE_KERNELS。这一点在 constants.py 第 49-55 行 中有直接体现DIFFUSERS_TRUST_REMOTE_KERNELS的计算条件是环境变量为真且未设置DIFFUSERS_DISABLE_REMOTE_CODE。调试检查Checks注意力调度器内置了若干调试检查用于在问题酿成大错前捕获常见错误。默认关闭可通过DIFFUSERS_ATTN_CHECKS环境变量开启export DIFFUSERS_ATTN_CHECKSyes开启后每次注意力运算前都会执行检查。检查主要分三类实现见 attention_dispatch.py 第 443-514 行设备检查device checks验证 query、key、value 张量位于同一设备_check_device部分后端还会额外要求位于 CUDA 设备或满足特定算力_check_device_cuda_atleast_smXY数据类型检查dtype checks确认 Q/K/V 的 dtype 一致且为 bfloat16 或 float16_check_qkv_dtype_bf16_or_fp16个别后端如 FA3进一步要求 bfloat16形状检查shape checks校验张量维度例如 query 与 key 的 head 维度一致、key 与 value 的序列长度一致、attention mask 与 key 的序列长度匹配并阻止同时混用 attention mask 与 causal 标志_check_attn_mask_or_causal当attn_mask不为空且is_causalTrue时直接抛错。开启检查后可用下面的片段直接验证某个后端在当前环境下是否可用import torch query torch.randn(1, 10, 8, 64, dtypetorch.bfloat16, devicecuda) # or mps, xpu, cpu key torch.randn(1, 10, 8, 64, dtypetorch.bfloat16, devicecuda) value torch.randn(1, 10, 8, 64, dtypetorch.bfloat16, devicecuda) 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})[!NOTE] 检查会给每一次注意力操作带来额外开销所以默认是关闭的建议在调试阶段开启正式跑推理时保持关闭以获得最佳性能。也可以直接配置注册表来开启检查from diffusers.models.attention_dispatch import _AttentionBackendRegistry _AttentionBackendRegistry._checks_enabled True需要注意dispatch_attention_fn内部开启检查时第 431-440 行还会自动剔除当前后端不支持的额外参数并给出 warning例如启用 GQA 时只有支持enable_gqa的后端才会收到该参数。Available backends 完整列表下表列出了当前版本可用的全部注意力后端及其变体与AttentionBackendName枚举一一对应见 attention_dispatch.py 第 215-255 行后端名称家族说明nativePyTorch native默认后端使用 PyTorch 的 scaled_dot_product_attentionflexFlexAttentionPyTorch FlexAttention 实现_native_cudnnPyTorch nativeCuDNN 优化注意力_native_efficientPyTorch native显存高效注意力_native_flashPyTorch nativePyTorch 的 FlashAttention_native_mathPyTorch native基于数学计算的注意力回退_native_npuPyTorch nativeNPU 优化注意力_native_xlaPyTorch nativeXLA 优化注意力flashFlashAttentionFlashAttention-2flash_hubFlashAttention来自 kernels 的 FlashAttention-2flash_varlenFlashAttention变长序列 FlashAttentionflash_varlen_hubFlashAttention来自 kernels 的变长 FlashAttentionaiter_fa2_hubAI Tensor Engine for ROCm来自 kernels 的 AMD ROCm FlashAttention-2flash_4_hubFlashAttentionFlashAttention-4_flash_3FlashAttentionFlashAttention-3_flash_varlen_3FlashAttention变长 FlashAttention-3_flash_3_hubFlashAttention来自 kernels 的 FlashAttention-3_flash_3_varlen_hubFlashAttention来自 kernels 的变长 FlashAttention-3sageSageAttention量化注意力INT8 QKsage_hubSageAttention来自 kernels 的量化注意力INT8 QKsage_varlenSageAttention变长 SageAttention_sage_qk_int8_pv_fp8_cudaSageAttentionINT8 QK FP8 PVCUDA_sage_qk_int8_pv_fp8_cuda_sm90SageAttentionINT8 QK FP8 PVSM90 架构_sage_qk_int8_pv_fp16_cudaSageAttentionINT8 QK FP16 PVCUDA_sage_qk_int8_pv_fp16_tritonSageAttentionINT8 QK FP16 PVTritonxformersxFormers显存高效注意力dispatch_attention_fn会以注册表记录的_supported_arg_names由函数签名自动推导为基准做参数过滤因此各后端只需实现自己支持的参数子集即可无缝接入见 attention_dispatch.py 第 278 行 与 第 438 行。如何选择与验证后端结合文档与源码给出如下实践建议追求零配置开箱即用优先选择native默认或 Hub 类后端_flash_3_hub、flash_hub、sage_hub等它们不需要本地编译 CUDA 扩展只需pip install kernels即可自动拉取预编译内核显存受限的大模型推理优先尝试 FlashAttention 系列flash/flash_varlen或 xFormersxformers两者均以显存效率著称SageAttention 的 int8 量化路径sage/_sage_*适合进一步降低显存与带宽开销AMD ROCm 平台使用aiter_fa2_hub获得针对 ROCm 优化的 FlashAttention-2对比与回归用attention_backend上下文管理器在同一个推理脚本里逐个切换后端对比速度与显存占用开启DIFFUSERS_ATTN_CHECKS便于在切换时第一时间捕获设备、dtype、形状等基础错误验证正确性仓库在 tests/models/test_attention_dispatch.py 中提供了注意力后端的相关测试如上下文并行下的test_attn_backend_backward_parity可参考其思路做不同后端输出的数值一致性对比。另外注意set_attention_backend与reset_attention_backend是ModelMixin的方法适用于任意继承它的模型如pipeline.transformer而attention_backend上下文管理器与dispatch_attention_fn均可从diffusers.models导入见 models/init.py 第 174 行 与 diffusers/init.py 第 369 行。若希望全局指定默认后端可在启动时设置DIFFUSERS_ATTN_BACKEND环境变量默认值native见 constants.py 第 45 行注册表的_active_backend会直接读取该变量作为初始值attention_dispatch.py 第 263 行。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考