19.8 现代 Attention API 与 Hugging Face Kernels

Author

jshn9515

Published

2026-06-19

Modified

2026-06-19

第 10 章我们已经完整讨论过 FlashAttention 的算法原理：

为什么标准 attention 会产生大量 HBM 读写？
如何通过分块计算避免物化完整 attention 矩阵？
Online Softmax 为什么能保证结果正确？

因此，这一节不再重新解释 FlashAttention。我们把重点转向一个更实际的问题：

今天在 PyTorch 和 Hugging Face 里写 attention，到底应该调用哪个 API？

现代深度学习框架通常不会要求用户直接调用某个特定版本的 FlashAttention kernel。更常见的使用方式是调用一个统一的 attention API，由框架根据输入条件和硬件自动选择合适的 kernel。

在 PyTorch 中，这个统一入口主要是 SDPA：

F.scaled_dot_product_attention(query, key, value)

围绕 SDPA，PyTorch 还提供：

torch.nn.attention；
SDPBackend；
sdpa_kernel；
FlexAttention。

PyTorch 之外，常见的高效 attention 生态还包括：

Meta 的 xFormers 库；
Hugging Face 的 kernels 库。

这一节的目标不是比较每个 kernel 的内部实现，而是建立一张选择地图：

普通 causal attention，优先使用 SDPA；
需要控制具体后端，使用 sdpa_kernel + SDPBackend；
需要自定义 mask、bias 或 attention 变体，使用 FlexAttention；
已有项目依赖 xFormers 或需要其 AttentionBias，继续使用 xFormers；
使用 Transformers 并希望动态加载社区 kernel，使用 Hugging Face kernels。

from functools import partial

import dnnlpy
import IPython.display as ipy
import kernels
import torch
import torch.nn as nn
import torch.nn.attention as attn
import torch.nn.attention.flex_attention as flex_attn
import torch.nn.functional as F
import xformers.ops as xops
from torch import Tensor
from transformers import AutoModelForCausalLM

print('PyTorch version:', torch.__version__)

device = dnnlpy.get_default_device()
print('Using device:', device)

19.8.1 SDPA：PyTorch 的统一 Attention 入口

Scaled Dot-Product Attention 的数学形式是：

\[ \operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_h}} + M \right)V \]

其中，\(M\) 可以表示 causal mask、padding mask 或其他 bias。

在 PyTorch 中，可以直接调用：

F.scaled_dot_product_attention(query, key, value)

我们可以先用一个示例验证 SDPA 的输入输出形状。

batch_size = 16
num_heads = 4
seq_len = 16
head_dim = 8

query = torch.randn(batch_size, num_heads, seq_len, head_dim)
key = torch.randn_like(query)
value = torch.randn_like(query)

output = F.scaled_dot_product_attention(query, key, value)

print('input shape:', query.shape)
print('output shape:', output.shape)

对于 GPT 的 causal self-attention，通常写成：

F.scaled_dot_product_attention(query, key, value, is_causal=True)

output = F.scaled_dot_product_attention(query, key, value, is_causal=True)

print('Causal attention output shape:', output.shape)

这里最重要的是：

SDPA 是 API 入口，不是某一个固定 kernel 的名字。

同一段 SDPA 代码，在不同硬件和输入条件下，可能使用不同后端。因此，不应该把调用 SDPA 等同于一定调用 FlashAttention。更准确的理解是：我描述需要哪种 attention，PyTorch 决定如何执行。

19.8.2 SDPA 为什么比手写 Attention 更适合作为默认实现

手写 Attention 往往是：

scores = query @ key.transpose(-2, -1)
scores = scores / math.sqrt(head_dim)
scores = scores.masked_fill(mask == 0, -inf)
probs = scores.softmax(dim=-1)
output = probs @ value

它很适合教学，但训练代码中直接这样写，需要显式创建一个 \((B, H, T, T)\) 大小的 scores 和 probs 张量。相反，SDPA 把整个 attention 作为一个高层操作交给框架：

Q, K, V
-> SDPA dispatcher
-> 选择可用 backend
-> 输出

这样做有几个好处。首先，SDPA 可以在内部使用 fused kernel，避免显式构造完整 attention 矩阵，从而节省显存和内存带宽。其次，SDPA 的后端可以随 PyTorch 更新，而不需要用户手动维护第三方 kernel。最后，SDPA 提供了 CPU 和 GPU 的 fallback 机制，使得模型代码不与某个第三方 kernel 强绑定。所以，真正写 attention 时，应该优先把手写 attention 替换为 SDPA。

这里还有一个容易忽略的细节。SDPA 的 dropout_p 会根据传入值执行 dropout，而不会自动读取模块的 training 状态（因为这是一个函数）。因此应该显式写：

dropout_p = self.dropout if self.training else 0.0

而不是在推理阶段仍然传入训练时的 dropout probability。

19.8.3 torch.nn.attention：控制 SDPA 行为的命名空间

torch.nn.attention 是围绕 SDPA 的控制与扩展命名空间，其中包括：

sdpa_kernel：一个上下文管理器，用于控制 SDPA 可以选择哪些后端；
SDPBackend：一个枚举类，表示当前 PyTorch 支持的 SDPA 后端；
FlexAttention：用于描述自定义 attention 变体的 API。

常见导入方式是：

from torch.nn.attention import SDPBackend, sdpa_kernel

当然，普通模型代码通常不需要固定后端。让 PyTorch 自动选择，往往更容易跨设备和跨版本运行。

19.8.3.1 SDPBackend：SDPA 有哪些后端

SDPBackend 是一个类似枚举的类，用于表示 SDPA 后端。常见成员包括：

MATH
FLASH_ATTENTION
EFFICIENT_ATTENTION
CUDNN_ATTENTION

不同成员对应不同算法。大体上可以理解为：

表 19.8.3.1 不同 SDPA 后端的对应关系
Backend	对应算法
`MATH`	使用普通 PyTorch / C++ 数学实现，兼容性最好
`FLASH_ATTENTION`	使用 FlashAttention 类算法，如 FA3 和 FA4
`EFFICIENT_ATTENTION`	使用 memory-efficient attention，来自 xFormers
`CUDNN_ATTENTION`	使用 NVIDIA cuDNN 提供的 attention 后端

这些名字描述的是后端类别，不代表所有设备都支持。例如 CPU 通常只支持 MATH，而 CUDA GPU 可能支持 Flash、Efficient 或 cuDNN。具体能否使用，取决于 dtype、shape、mask 和硬件。

19.8.3.2 用 sdpa_kernel 强制选择后端

sdpa_kernel 是一个 context manager，用于控制 SDPA 可以选择哪些后端。

例如，强制使用数学后端：

with attn.sdpa_kernel(attn.SDPBackend.MATH):
    output = F.scaled_dot_product_attention(query, key, value, is_causal=True)

print('Attention output shape:', output.shape)

也可以传入多个候选后端：

with attn.sdpa_kernel(
    [
        attn.SDPBackend.MATH,
        attn.SDPBackend.FLASH_ATTENTION,
        attn.SDPBackend.EFFICIENT_ATTENTION,
        attn.SDPBackend.CUDNN_ATTENTION,
    ]
):
    output = F.scaled_dot_product_attention(query, key, value, is_causal=True)

print('Attention output shape:', output.shape)

这样表示允许使用这些 backend，PyTorch 会根据输入条件和硬件选择最优的可用后端。如果设置 set_priority=True，则会按照传入顺序选择可用后端。如果输入不满足要求，程序可能报错或给出无法使用该 backend 的提示。

因此，下面这种写法更适合 benchmark，而不是通用模型代码：

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    F.scaled_dot_product_attention(query, key, value, is_causal=True)

日常训练通常更推荐不写 sdpa_kernel，让 PyTorch dispatcher 自动选择后端。

19.8.3.3 如何正确比较不同 SDPA 后端

比较不同 attntion backend 时，不能只运行一次然后用 time.time()。因为 GPU kernel 通常是异步执行的，而且第一次运行可能包含初始化和编译开销。

一个更合理的 benchmark 流程是：

先 warmup，确保 kernel 已经编译；
在计时前调用 synchronize()，确保前面的 kernel 执行完毕；
重复多次执行 kernel；
再次调用 synchronize()，确保所有 kernel 执行完毕；
计算平均时间。

def benchmark_sdpa(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_backend: attn.SDPBackend,
    num_warmup: int = 5,
    num_steps: int = 20,
) -> float:
    if device.type not in {'cpu', 'cuda', 'xpu', 'mps'}:
        raise AssertionError('Benchmark only supports CPU, CUDA, XPU, MPS devices.')

    with attn.sdpa_kernel(attn_backend):
        for _ in range(num_warmup):
            F.scaled_dot_product_attention(query, key, value, is_causal=True)

        backend = getattr(torch, device.type)
        backend.synchronize()

        start = torch.Event(enable_timing=True)
        end = torch.Event(enable_timing=True)

        start.record()

        for _ in range(num_steps):
            F.scaled_dot_product_attention(
                query,
                key,
                value,
                is_causal=True,
            )

        end.record()
        backend.synchronize()

    return start.elapsed_time(end) / num_steps

接下来我们创建几个比较大的张量，测试不同 backend 的平均执行时间。

batch_size = 64
num_heads = 32
seq_len = 1024
head_dim = 64

query = torch.randn(
    batch_size,
    num_heads,
    seq_len,
    head_dim,
    dtype=torch.bfloat16,
    xdevice=device,
)
key = torch.randn_like(query)
value = torch.randn_like(query)

for backend in [
    attn.SDPBackend.MATH,
    attn.SDPBackend.FLASH_ATTENTION,
    attn.SDPBackend.EFFICIENT_ATTENTION,
    attn.SDPBackend.CUDNN_ATTENTION,
]:
    try:
        avg_time = benchmark_sdpa(query, key, value, backend)
        print(f'Backend: {backend.name}, Average time: {avg_time:.2f} ms.')
    except RuntimeError as err:
        print(f'Backend: {backend.name}, Error: {err}.')

Backend: MATH, Average time: 438.57 ms.
Backend: FLASH_ATTENTION, Average time: 7.17 ms.
Backend: EFFICIENT_ATTENTION, Average time: 11.58 ms.
Backend: CUDNN_ATTENTION, Average time: 7.64 ms.

当然，训练场景中只测 forward 不够，还应该测 forward + backward，并关注 peak memory 和完整训练 step 的时间。因为一个 kernel 的 forward 更快，不代表整个训练 step 一定更快。

19.8.4 FlexAttention：当 SDPA 的 Mask 不够灵活

FlexAttention (Dong et al. 2024) 是 PyTorch 2.5 引入的一个新特性，目标是让用户可以用 Python 函数描述自定义 attention 语义，然后由 PyTorch 编译成高效 fused kernel。

Tip

PyTorch 2.5：首次引入 FlexAttention，支持 CUDA；
PyTorch 2.6：增加对 x86 CPU 的支持；
PyTorch 2.7：进一步完善 FlexAttention API；
PyTorch 2.9：增加对 XPU 的支持；
PyTorch 2.13：增加对 MPS 的支持。

SDPA 适合常见的 attention 语义，例如：

Casual mask：GPT 的自回归 causal attention；
Padding mask：处理变长序列时的 padding；
普通 attention bias：例如 ALiBi 等。

但现代模型会出现更复杂的 attention 模式，例如：

Sliding Window Attention：局部窗口的 mask；
Document Masking：文档级别的 mask；
Prefix Language Model：前缀语言模型的 mask；
Local + Global Attention：局部和全局 attention 组合；
Block Sparse Attention：块稀疏 attention。

如果每一种变体都单独手写 CUDA 或 Triton kernel，维护成本会非常高。

因此，FlexAttention 的目标是：

让用户用普通 Python 函数描述 score 或 mask 的修改规则，再由 PyTorch 编译成高效 fused kernel。

其核心 API 位于：

torch.nn.attention.flex_attention

FlexAttention 常见的两个概念是 score_mod 和 mask_mod。

score_mod 修改 softmax 之前的 attention score。例如，给 score 加一个与相对位置有关的 bias：

def score_mod(
    score: Tensor,
    batch_idx: IntTensor,
    head_idx: IntTensor,
    q_idx: IntTensor,
    kv_idx: IntTensor,
) -> Tensor: ...

mask_mod 则返回某个 query-key 位置是否允许参与 attention：

def causal_mask(
    batch: Tensor,
    head_idx: IntTensor,
    q_idx: IntTensor,
    kv_idx: IntTensor,
) -> BoolTensor: ...

我们写个小例子来演示一下。首先，定义 causal mask：

def causal_mask(
    batch: Tensor, head_idx: Tensor, q_idx: Tensor, kv_idx: Tensor
) -> Tensor:
    return q_idx >= kv_idx

然后使用 create_block_mask() 创建 block-sparse mask 表示：

batch_size = 16
num_heads = 4
q_len = 128
kv_len = 128
head_dim = 32

query = torch.randn(batch_size, num_heads, q_len, head_dim, device=device)
key = torch.randn_like(query)
value = torch.randn_like(query)

block_mask = flex_attn.create_block_mask(
    causal_mask,
    B=batch_size,
    H=num_heads,
    Q_LEN=q_len,
    KV_LEN=kv_len,
    device=device,
)
opts = {'BACKEND': 'AUTO'}  # Add more options if needed
flex_attention = partial(flex_attn.flex_attention, kernel_options=opts)
flex_attention = torch.compile(flex_attention, dynamic=False)

output = flex_attention(query, key, value, block_mask=block_mask)
print('FlexAttention output shape:', output.shape)

真实使用中，FlexAttention 通常与 torch.compile 配合，以生成针对当前 attention variant 的高效实现。这尤其适合 attention 语义特殊，但又希望得到 fused kernel 的场景。

需要注意，FlexAttention 的 API 和支持范围仍在快速演进。请以最新的 PyTorch 文档为准。

19.8.5 xFormers：PyTorch 原生 SDPA 之前的重要高效算子库

xFormers (Lefaudeux et al. 2022) 是 Meta 维护的 PyTorch 扩展库，提供可组合的 Transformer 组件和优化算子。与 SDPA 不同，xFormers 并不是 PyTorch 官方的一部分，而是一个第三方库。

xFormers 最常见的 API 是：

xformers.ops.memory_efficient_attention

典型输入 layout 是：

\[ (B, T, H, d_h) \]

这与 SDPA 常见的：

\[ (B, H, T, d_h) \]

不同，使用时要特别注意。

简化代码如下：

output = xops.memory_efficient_attention(
    query, key, value, attn_bias=xops.LowerTriangularMask()
)

print('xFormers output shape:', output.shape)

xFormers 的一个重要特点是提供 Attention Bias 抽象，例如：

LowerTriangularMask
LowerTriangularFromBottomRightMask
LowerTriangularMaskWithTensorBias
BlockDiagonalMask
BlockDiagonalCausalMask

这些 bias 可以避免显式创建完整的 dense mask。

前几年，xFormers 是 PyTorch 官方 SDPA 之前的主要高效 attention 算子库。随着 PyTorch 逐渐完善 FlexAttention 和 SDPA，并且提供了 EFFICIENT_ATTENTION 后端，xFormers 的必要性有所下降。当前，xFormers 更多用于 PyTorch 尚未完全覆盖或需要更细粒度控制的场景，例如特殊的 attention bias、paged attention、长上下文和解码相关算子上。

19.8.6 Hugging Face kernels：从 Hub 动态加载 kernel

Hugging Face 的 kernels 库 (Hugging Face 2026) 解决的问题和 SDPA 不完全一样。

SDPA 是 PyTorch 内部统一 attention API，而 kernels 则是一个 kernel 的分发、版本管理和动态加载系统。所有的算子都可以作为 Hugging Face Hub 上的一类仓库发布。这类算子通常是预编译好的，用户通过 kernels 库直接加载与当前环境兼容的实现，避免了本地编译和依赖冲突。

基本形式是：

module = kernels.get_kernel('kernels-community/flash-attn3', version=1)

然后从返回的模块中取得具体函数。

Hugging Face Transformers 也支持在加载模型时指定 attention kernel：

model = AutoModelForCausalLM.from_pretrained(
    'Qwen/Qwen3-0.6B',
    attn_implementation=('kernels-community/flash-attn3'),
    device_map='auto',
)
ipy.clear_output()

这里的 attn_implementation 参数表示希望使用的 kernel。Transformers 会尝试从 Hugging Face Hub 下载并加载该 kernel。如果本地已经有缓存，则直接使用缓存。

除了 attn_implementation，Transformers 还提供了一个全局开关 use_kernels，当设置为 True 时，Transformers 会尝试为所有支持的层加载匹配的 kernel，而不仅限于 attention。

model = AutoModelForCausalLM.from_pretrained(
    'Qwen/Qwen3-0.6B',
    attn_implementation=('kernels-community/flash-attn3'),
    device_map='auto',
    use_kernels=True,
)
ipy.clear_output()

这种设计的优势在于，kernel 可以独立于 Transformers 发布，可以针对不同硬件提供不同实现，并且可以固定 kernel 版本，从而减少本地编译和依赖冲突。但它也意味着加载 kernel 时会执行任意代码，因此只能加载可信来源。如果确实需要加载第三方的 kernel，需要显式设置 allow_all_kernels=True。

19.8.7 这些方案应该怎么选

可以先使用下面的决策顺序。

如果是自己实现普通 GPT，优先使用：

import torch.nn.functional as F

F.scaled_dot_product_attention(query, key, value, is_causal=True)

让 PyTorch 自动选择 backend。

如果是想确认或比较不同 SDPA 后端，可以使用：

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    F.scaled_dot_product_attention(query, key, value, is_causal=True)

只在 benchmark 或调试区域强制 backend。

如果需要复杂 attention 变体，例如 sliding window、document mask、soft-capping 或 block sparse pattern，可以考虑使用 FlexAttention：

from torch.nn.attention.flex_attention import flex_attention

flex_attention = torch.compile(flex_attention, dynamic=False)
flex_attention(query, key, value, block_mask=...)

如果项目已经依赖 xFormers，或者需要其 AttentionBias 抽象，可以继续使用 xFormers：

import xformers.ops as xops

xops.memory_efficient_attention(query, key, value, attn_bias=...)

而如果希望从 Hugging Face Hub 动态加载 kernel，或者希望固定 kernel 版本，可以使用 Hugging Face kernels：

import kernels

module = kernels.get_kernel('kernels-community/flash-attn3', version=1)
module.flash_attention(query, key, value)

最重要的是，不要仅根据名称判断哪个一定最快。最终仍然应该测量完整训练 step 的时间、tokens per second、peak memory、数值稳定性和硬件兼容性。

19.8.8 本章小结

这一节，我们整理了现代 attention 的算子生态。

对于 PyTorch，最重要的入口是：

torch.nn.functional.scaled_dot_product_attention

SDPA 描述 attention 语义，并由 PyTorch 根据输入条件和硬件自动选择最优后端。

如果需要控制后端，可以使用 SDPBackend 和 sdpa_kernel，常见后端包括：

MATH
FLASH_ATTENTION
EFFICIENT_ATTENTION
CUDNN_ATTENTION

如果需要实现复杂 attention 变体，可以使用 FlexAttention，通过 score_mod 和 mask_mod 描述自定义规则，再交给 PyTorch 和 Triton 进行编译。

xFormers 仍然是重要的优化算子库，虽然很多功能已经被 PyTorch 的 SDPA 和 FlexAttention 覆盖，但它仍然提供了丰富的 AttentionBias 抽象，尤其适合已有项目依赖、diffusers 生态以及需要结构化 AttentionBias 的场景。

Hugging Face kernels 则进一步解决了 kernel 的分发、版本管理和动态加载问题。它允许 Transformers 在加载模型时指定 kernel，并从 Hub 下载预编译的实现，避免本地编译和依赖冲突。

可以把整张地图压缩成：

普通 attention，优先使用 SDPA；
需要控制后端，使用 SDPBackend + sdpa_kernel
需要自定义 mask、bias 或 attention 变体，使用 FlexAttention；
已有项目依赖 xFormers 或需要其 AttentionBias，继续使用 xFormers；
使用 Transformers 并希望动态加载社区 kernel，使用 Hugging Face kernels。

当然，对于大多数自己实现的 GPT，最合理的起点仍然是：

F.scaled_dot_product_attention(...)

下一节，我们将进入单机多卡训练，讨论 DDP、ZeRO 与 FSDP。单卡算子优化解决的是不同算子在单个 GPU 内部如何高效执行，而多卡训练则要回答：

如何把模型计算和训练状态分摊到多张 GPU。

References

Dong, Juechu, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. 2024. Flex Attention: A Programming Model for Generating Optimized Attention Kernels. https://arxiv.org/abs/2412.05496.

Hugging Face. 2026. Kernels: Build Compute Kernels and Load Them from the Hub. Released. https://github.com/huggingface/kernels.

Lefaudeux, Benjamin, Francisco Massa, Diana Liskovich, et al. 2022. xFormers: A Modular and Hackable Transformer Modelling Library. Https://github.com/facebookresearch/xformers, released.

Reuse

CC BY-NC 4.0