Chaos and Order

💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

1. 引言
2. Standard Attention 的内存问题
3. GPU 内存层级结构
4. IO Complexity 分析
5. Tiling 技巧：以适配 SRAM 的块为单位运算
6. Online Softmax（Safe Softmax）算法
7. Backward Pass 的 Recomputation 策略
8. FlashAttention-2 的改进
9. FlashAttention-3 的最新进展
10. 基准测试：速度/内存对比
11. 与 PyTorch torch.nn.functional.scaled_dot_product_attention 的集成
12. 总结与核心要点
References

1. 引言

作为 Transformer 架构核心的 Self-Attention，负责计算序列内所有 token 对之间的关系。这一运算具有强大的表达能力，但存在一个根本性的局限：相对于序列长度 $N$ ，时间与内存复杂度都以 $O(N^2)$ 增长。GPT-4、LLaMA、Gemini 等最新 LLM 若要处理 128K 以上的长上下文，就必须切实解决这一 $O(N^2)$ 瓶颈。

FlashAttention（Dao et al., 2022）解决这一问题的方式是：不做任何近似。其核心思路简单却深刻——并非降低 attention 运算本身的计算量，而是最小化 GPU 内存层级之间的数据搬移（IO）。本文将从 GPU 硬件的视角系统分析 FlashAttention 的原理，并梳理其向 FlashAttention-2、FlashAttention-3 演进的历程。

2. Standard Attention 的内存问题

2.1 Standard Attention 的计算流程

Standard Self-Attention 按如下方式计算。给定输入 $Q, K, V \in \mathbb{R}^{N \times d}$ ：

S = QK^T \in \mathbb{R}^{N \times N}

P = \text{softmax}(S) \in \mathbb{R}^{N \times N}

O = PV \in \mathbb{R}^{N \times d}

其中 $N$ 是序列长度， $d$ 是 head dimension。

2.2 内存复杂度分析

问题的核心在于中间矩阵 $S$ 和 $P$ 。这些矩阵的大小为 $N \times N$ ，相对于序列长度所需的内存是二次的（quadratic）。具体数值如下：

序列长度 ( $N$ )	Attention 矩阵大小	FP16 内存
1,024	1M 个元素	2 MB
4,096	16.7M 个元素	33 MB
16,384	268M 个元素	536 MB
65,536	4.3B 个元素	8.6 GB
131,072	17.2B 个元素	34.4 GB

这些数值针对的是：单个 head、单个 batch。在 multi-head attention 中，乘以 head 数量 $h$ 、再乘以 batch size $B$ ，实际内存占用会大得多。在序列长度为 65,536 时，仅单个 head 就已经会消耗 A100 80GB GPU HBM 的相当一部分。

2.3 HBM 瓶颈

在 standard attention 的实现中，这个 $N \times N$ 矩阵会被具体化（materialize）到 GPU HBM（High Bandwidth Memory）中。也就是说，先计算 $S = QK^T$ 并写入 HBM，再读出用于 softmax，将结果 $P$ 写入 HBM，然后再读出用于计算 $O = PV$ 。这一过程中对 HBM 的读写次数为 $\Omega(Nd + N^2)$ 。

实际 GPU 上这一运算之所以缓慢，原因在于：瓶颈不是计算（compute），而是内存访问（memory access）。A100 GPU 的计算吞吐量为 312 TFLOPS（FP16），而 HBM 带宽却只有约 2 TB/s。Attention 运算的 arithmetic intensity（计算量与内存访问量之比）较低，是典型的 memory-bound 运算。

3. GPU 内存层级结构

要理解 FlashAttention，就必须准确了解 GPU 的内存层级。

3.1 HBM（High Bandwidth Memory）

容量：A100 上为 40GB 或 80GB
带宽：约 1.5-2.0 TB/s（A100 80GB SXM：2,039 GB/s）
访问延迟：约 200-600 个周期
角色：GPU 的主内存，模型参数、输入张量、输出张量等所有数据都存储于此

3.2 SRAM（On-chip Shared Memory）

容量：A100 上每个 SM 约 192KB，共计约 20MB（108 个 SM）
带宽：约 19 TB/s
访问延迟：约 20-30 个周期
角色：每个 Streaming Multiprocessor（SM）内部的高速片上内存

3.3 核心的非对称性

SRAM 与 HBM 之间存在着显著的非对称性：

特性	SRAM	HBM
带宽	~19 TB/s	~2 TB/s
容量	~20 MB	40-80 GB
访问延迟	20-30 个周期	200-600 个周期

SRAM 比 HBM 快约 10 倍，但容量却小约 4000 倍。 FlashAttention 的核心洞察，正是积极利用这一非对称性：与其将整个 $N \times N$ 矩阵在 HBM 中具体化，不如以能装进 SRAM 的小块为单位执行运算，从而大幅减少 HBM 访问。

4. IO Complexity 分析

4.1 Standard Attention 的 IO Complexity

Standard attention 表现出如下的 HBM 访问模式：

从 HBM 读取 $Q, K$ ，计算 $S = QK^T$ ，将 $S$ 写入 HBM： $\Theta(Nd + N^2)$ IO
从 HBM 读取 $S$ ，计算 $P = \text{softmax}(S)$ ，将 $P$ 写入 HBM： $\Theta(N^2)$ IO
从 HBM 读取 $P, V$ ，计算 $O = PV$ ，将 $O$ 写入 HBM： $\Theta(Nd + N^2)$ IO

HBM 总访问量： $\Theta(Nd + N^2)$

由于序列长度 $N$ 通常远大于 head dimension $d$ （一般为 64 或 128）， $N^2$ 项占主导地位。

4.2 FlashAttention 的 IO Complexity

FlashAttention 通过 tiling 将 HBM 访问量降低到：

O\left(\frac{N^2 d^2}{M}\right)

其中 $M$ 是 SRAM 大小。直观上，SRAM 越大，一次能处理的块就越大，HBM 访问也随之减少。

4.3 最优性证明（Lower Bound）

论文进一步证明了以下下界（lower bound）：

定理：对于满足 $d \leq M \leq Nd$ 的任意 SRAM 大小 $M$ ，任何计算 exact attention 的算法都需要 $\Omega(N^2 d^2 / M)$ 次 HBM 访问。

这意味着 FlashAttention 在 IO complexity 意义上是最优的（optimal）。除去常数因子和多项式对数因子，不可能用更少的 HBM 访问来计算 exact attention。

4.4 数值示例

在 A100 上，SRAM 大小 $M \approx 192$ KB、head dimension $d = 64$ 、序列长度 $N = 4096$ 时：

Standard attention IO： $\Theta(Nd + N^2) \approx 4096 \times 64 + 4096^2 \approx 17M$ 个元素
FlashAttention IO： $\Theta(N^2 d^2 / M) \approx 4096^2 \times 64^2 / (192 \times 512) \approx 7M$ 个元素（随块大小而变化）

实际上，由于 $N^2$ 大小的中间矩阵完全不会被写入 HBM，节省效果会更大。序列长度越长，这一效果就越显著。

5. Tiling 技巧：以适配 SRAM 的块为单位运算

5.1 算法概览

FlashAttention 的核心算法如下：

把 $Q$ 分成 $T_r = \lceil N / B_r \rceil$ 个块： $Q_1, Q_2, \ldots, Q_{T_r}$ ，每块大小为 $B_r \times d$
把 $K, V$ 分成 $T_c = \lceil N / B_c \rceil$ 个块： $K_1, \ldots, K_{T_c}$ 及 $V_1, \ldots, V_{T_c}$ ，每块大小为 $B_c \times d$
块大小 $B_r, B_c$ 根据 SRAM 大小 $M$ 设定： $B_c = \lceil M / (4d) \rceil$ ， $B_r = \min(\lceil M / (4d) \rceil, d)$

5.2 Forward Pass 伪代码

Algorithm: FlashAttention Forward Pass
---------------------------------------
Input: Q, K, V in HBM, SRAM size M
Output: O in HBM

1. 设置块大小：B_c = ceil(M / 4d), B_r = min(ceil(M / 4d), d)
2. 在 HBM 中初始化 O = zeros(N, d), l = zeros(N), m = -inf * ones(N)

3. for j = 1 to T_c:                        # 外层循环：K, V 块
     将 K_j, V_j 从 HBM 加载到 SRAM

     for i = 1 to T_r:                      # 内层循环：Q 块
       将 Q_i, O_i, l_i, m_i 从 HBM 加载到 SRAM

       # 在 SRAM 中执行块级运算
       S_ij = Q_i @ K_j^T                   # (B_r x B_c)
       m_ij = rowmax(S_ij)
       P_ij = exp(S_ij - m_ij)
       l_ij = rowsum(P_ij)

       # 与之前块的统计量结合（Online Softmax）
       m_new = max(m_i, m_ij)
       l_new = exp(m_i - m_new) * l_i + exp(m_ij - m_new) * l_ij

       # 更新输出（含 rescaling）
       O_i = diag(exp(m_i - m_new))^(-1) * (diag(l_i) * O_i)
             + diag(exp(m_ij - m_new))^(-1) * P_ij @ V_j
       O_i = diag(l_new)^(-1) * O_i

       # 更新统计量
       m_i = m_new, l_i = l_new

       将 O_i, l_i, m_i 写回 HBM
     end for
   end for

4. return O

5.3 为什么这样可行

关键在于， $N \times N$ 大小的 attention 矩阵 $S$ 和 $P$ 完全不会在 HBM 中具体化。每个 $B_r \times B_c$ 大小的块 $S_{ij}$ 都在 SRAM 内计算，随即用于更新 softmax 统计量和累加输出，之后即被丢弃。

使这一切成为可能的数学技巧，正是 Online Softmax。

6. Online Softmax（Safe Softmax）算法

6.1 Standard Softmax 的问题

Softmax 是一种全局运算（global operation）。对于行向量 $x = [x_1, \ldots, x_N]$ ：

\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_{j=1}^{N} e^{x_j}}

要计算它，必须一次性看到整行才能求出分母的总和。这正是让 tiling 变得困难的根本障碍——仅凭块 $S_{i1}$ 无法完成 softmax 的计算，因为分母会随其余块 $S_{i2}, S_{i3}, \ldots$ 的值而变化。

另外，出于数值稳定性的考虑，会使用「safe softmax」：

\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_{j=1}^{N} e^{x_j - m}}, \quad m = \max_j x_j

这同样需要全局最大值 $m$ ，因此必须先扫描整行。

6.2 Online Softmax 技巧

Online Softmax（Milakov & Gimelshein, 2018）的核心思路是：在维持 running statistics 的同时，以块为单位递增（incremental）地计算 softmax。

每一行都维护两个标量：

$m$ ：目前为止所见元素的最大值（running max）
$\ell$ ：目前为止的归一化常数（running sum of exponentials）

当新的块 $S_{ij}$ 到来时：

计算新块的逐行最大值： $\tilde{m} = \text{rowmax}(S_{ij})$
更新全局最大值： $m_{\text{new}} = \max(m, \tilde{m})$
对之前的归一化常数做 rescale： $\ell_{\text{new}} = e^{m - m_{\text{new}}} \cdot \ell + e^{\tilde{m} - m_{\text{new}}} \cdot \tilde{\ell}$
对之前的输出也做 rescale： $O_{\text{new}} = \frac{e^{m - m_{\text{new}}} \cdot \ell \cdot O + e^{\tilde{m} - m_{\text{new}}} \cdot \tilde{P}V_j}{\ell_{\text{new}}}$

这一过程从数学角度看是精确的（exact），而非近似。无论以何种顺序处理各个块，最终结果都与 standard attention 逐位相同（除去浮点运算顺序导致的微小数值差异）。

6.3 数学正当性

证明的核心在于 softmax 的 rescaling property：

\frac{e^{x_i - m'}}{\sum_j e^{x_j - m'}} = \frac{e^{x_i - m} \cdot e^{m - m'}}{\sum_j e^{x_j - m} \cdot e^{m - m'}} = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}

即便最大值从 $m$ 更新为 $m'$ ，由于分子和分母都乘上了相同的因子 $e^{m - m'}$ ，比值也不会改变。正是这一性质，使得之前块的结果可以安全地按新的最大值重新 rescale。

7. Backward Pass 的 Recomputation 策略

7.1 Standard Backward Pass 的问题

在 standard attention 的 backward pass 中，梯度计算需要用到 forward pass 中保存的中间矩阵 $S$ 和 $P$ 。由于它们的大小为 $N \times N$ ，在 forward 中保存、再于 backward 中读回，需要 $O(N^2)$ 内存。

7.2 FlashAttention 的 Recomputation

FlashAttention 采用了 gradient checkpointing 的一种变体。在 forward pass 中不保存 $S$ 和 $P$ ，取而代之只保存：

最终输出 $O \in \mathbb{R}^{N \times d}$
Softmax 归一化统计量 $m, \ell \in \mathbb{R}^{N}$ （逐行的最大值与总和）

在 backward pass 中，利用这些统计量以及原始的 $Q, K, V$ ，在 SRAM 中重新计算（recompute）所需的 $S$ 和 $P$ 块。这一 recomputation 会带来额外的 FLOP，但能大幅减少 HBM 访问。

7.3 Recomputation 的悖论效应

一般来说，gradient checkpointing 是以速度换内存。但 FlashAttention 的 recomputation 反而连速度也一并提升。原因如下：

FLOP 会增加：forward 中计算过一次的内容，在 backward 中又重新计算了一遍，因此总 FLOP 略有增加。
HBM IO 会减少：把 $N \times N$ 大小的 $S$ 、 $P$ 写入并读出 HBM 的开销消失了。

在现代 GPU 上，HBM 访问远比计算慢，因此 IO 减少带来的收益要大于 FLOP 增加带来的损失。实验结果显示，recomputation 带来的额外运行时开销低于 5%，而内存占用则从 $O(N^2)$ 降到了 $O(N)$ 。

7.4 内存节省效果

序列长度	Standard Attention 内存	FlashAttention 内存	节省比例
1K	~2 MB	~0.13 MB	~15x
2K	~8 MB	~0.26 MB	~30x
4K	~33 MB	~0.52 MB	~63x
8K	~131 MB	~1.04 MB	~126x

得益于这一节省效果，在相同的 GPU 内存下可以处理更长的序列，或使用更大的 batch size。

8. FlashAttention-2 的改进

Dao（2023）在 FlashAttention-2 中引入了三项核心改进。

8.1 最小化 Non-matmul FLOP

A100 GPU 的 Tensor Core 在矩阵乘法（matmul）上提供 312 TFLOPS（FP16），但 non-matmul 运算（softmax 的 exp、max、sum 等）只有 19.5 TFLOPS（FP32），慢了约 16 倍。在 FlashAttention-1 中，non-matmul 运算所占的比重相当可观。

FlashAttention-2 通过重新组织算法来最小化这些 non-matmul FLOP。具体来说，减少了 rescaling 运算的次数，并更高效地执行 softmax 统计量更新。核心改动在于，把最终的 rescaling 改成只在循环末尾执行一次。

8.2 并行性提升：Sequence Length 维度的并行化

FlashAttention-1 只在 batch 维度和 head 维度上做并行化。当 batch size 较小或 head 数较少时，无法充分利用 GPU 的 SM（Streaming Multiprocessor）。

FlashAttention-2 在 sequence length 维度上也做了并行化。做法是把外层循环改成以 Q 块为基准（外层循环不再是 $K, V$ 块，而是 $Q$ 块），使每个 Q 块都能在独立的 thread block 中处理。这一改动大幅提升了 forward pass 的 occupancy。

8.3 Work Partitioning 优化

Thread block 内部各 warp 之间的工作分配也得到了改进：

FlashAttention-1：把 K, V 分给 4 个 warp，各 warp 独立计算 $QK^T$ 后再同步结果。这种方式会产生通过 shared memory 通信和同步的开销。
FlashAttention-2：把 Q 分给 4 个 warp，K 和 V 则由所有 warp 共享。每个 warp 独立计算 Q 中不同部分对应的输出，因此，无需 warp 间通信。

8.4 性能结果

综合这三项改进：

相较 FlashAttention-1，约 2 倍加速
在 A100 上，FP16/BF16 达到 230 TFLOPS（约为理论峰值的 73%）
相较 standard PyTorch attention，最高 9 倍加速
逼近 GEMM（矩阵乘法）运算的效率

9. FlashAttention-3 的最新进展

FlashAttention-3（Shah et al., 2024）借助 NVIDIA Hopper 架构（H100）的新硬件特性，又向前迈进了一步。

9.1 Hopper GPU 的新特性

相比 A100，H100 GPU 提供以下核心特性：

WGMMA（Warpgroup Matrix Multiply-Accumulate）：吞吐量远高于 A100 的 mma.sync 的新 Tensor Core 指令
TMA（Tensor Memory Accelerator）：专门负责 global memory 与 shared memory 之间数据传输的硬件单元，索引计算和边界检查都由硬件处理

9.2 三项核心技巧

1. 通过 Warp Specialization 实现异步执行

将计算（WGMMA）与数据搬移（TMA）分配给不同的 warp group，以流水线方式重叠（overlap）执行。一个 warp group 在计算当前块的同时，另一个 warp group 会 prefetch 下一个块的数据。

2. Matmul 与 Softmax 的 Interleaving

以往的做法是 matmul 结束后再执行 softmax，然后再执行 matmul，是一种顺序方式。FlashAttention-3 将其交织（interleave）起来，让 matmul 和 softmax 在不同的硬件单元上同时执行。Tensor Core 计算下一个块的 $QK^T$ 的同时，CUDA Core 处理当前块的 softmax。

3. FP8 低精度支持

利用 H100 的 FP8 Tensor Core，将吞吐量提高 2 倍。单纯做 FP8 量化会降低精度，但 FlashAttention-3 用两项技巧解决了这一问题：

Block quantization：以块为单位维护各自独立的 scale factor，以保留动态范围
Incoherent processing：先乘以一个随机正交矩阵将 outlier 分散开，再进行量化，从而相较 FP8 baseline 实现 2.6 倍更低的数值误差

9.3 性能结果

FlashAttention-3 在 H100 上的性能：

配置	TFLOPS	GPU 利用率
FP16 FlashAttention-2	~400	~50%
FP16 FlashAttention-3	~740	~75%
FP8 FlashAttention-3	~1,200	~75%

在 FP16 下相较 FlashAttention-2 实现了 1.5-2.0 倍加速，FP8 下则达到了接近 1.2 PFLOPS 的性能。

10. 基准测试：速度/内存对比

10.1 Attention Forward Pass 速度（A100 80GB，FP16）

FlashAttention 论文及后续基准测试报告的主要数值如下：

序列长度	Standard Attention	FlashAttention	FlashAttention-2	Speedup（FA2 vs Std）
512	12.2 ms	3.5 ms	1.9 ms	6.4x
1K	45.8 ms	7.8 ms	4.1 ms	11.2x
2K	178 ms	18.9 ms	9.8 ms	18.2x
4K	710 ms	52.3 ms	27.1 ms	26.2x
8K	OOM	145 ms	75 ms	-
16K	OOM	520 ms	270 ms	-

序列长度越长，加速效果就越显著。在 8K 及以上，standard attention 会因 OOM（Out of Memory）而无法执行，FlashAttention 则可以毫无问题地处理。

10.2 端到端训练性能

模型	Standard	FlashAttention	Speedup
BERT-large（seq 512）	100%（MLPerf 基准）	115%	1.15x
GPT-2（seq 1K）	100%	300%	3.0x
Long-range Arena（seq 1K-4K）	100%	240%	2.4x

10.3 内存使用量对比

FlashAttention 的 attention 运算内存相对序列长度呈线性增长（linear），与 standard attention 呈二次增长（quadratic）相比有着显著改善：

序列长度 2K：内存节省约 10 倍
序列长度 4K：内存节省约 20 倍
序列长度 64K：standard attention 在 A100 80GB 上也会 OOM，FlashAttention 则可正常运行

11. 与 PyTorch `torch.nn.functional.scaled_dot_product_attention` 的集成

11.1 原生集成

从 PyTorch 2.0 开始，FlashAttention 已原生集成进 torch.nn.functional.scaled_dot_product_attention（SDPA）。从 PyTorch 2.2 开始，FlashAttention-2 成为默认 backend。

import torch
import torch.nn.functional as F

# 基本用法 - 自动选择 FlashAttention backend
query = torch.randn(batch_size, num_heads, seq_len, head_dim,
                    device='cuda', dtype=torch.float16)
key = torch.randn(batch_size, num_heads, seq_len, head_dim,
                  device='cuda', dtype=torch.float16)
value = torch.randn(batch_size, num_heads, seq_len, head_dim,
                    device='cuda', dtype=torch.float16)

# PyTorch 会自动选择最优的 backend
output = F.scaled_dot_product_attention(query, key, value)

11.2 显式指定 Backend

可以强制使用或排除特定的 backend：

from torch.nn.attention import sdpa_kernel, SDPBackend

# 仅使用 FlashAttention backend
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output = F.scaled_dot_product_attention(query, key, value)

# 仅使用 Memory-efficient attention backend
with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
    output = F.scaled_dot_product_attention(query, key, value)

# 使用 Math（naive）backend - 用于调试
with sdpa_kernel(SDPBackend.MATH):
    output = F.scaled_dot_product_attention(query, key, value)

# 使用 CuDNN backend（PyTorch 2.2+）
with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
    output = F.scaled_dot_product_attention(query, key, value)

11.3 与 Causal Mask 一起使用

LLM 自回归生成中不可或缺的 causal mask 也同样得到支持：

# 通过 is_causal=True 应用 causal mask
# FlashAttention 会在 fused kernel 内部处理 causal mask，无需额外内存
output = F.scaled_dot_product_attention(
    query, key, value,
    is_causal=True
)

# 使用自定义 attention mask
attn_mask = torch.tril(torch.ones(seq_len, seq_len, device='cuda', dtype=torch.bool))
output = F.scaled_dot_product_attention(
    query, key, value,
    attn_mask=attn_mask
)

11.4 Backend 选择条件

PyTorch SDPA 选择 FlashAttention backend 需要满足以下条件：

dtype：float16 或 bfloat16（不支持 float32）
device：CUDA GPU（不支持 CPU）
head dimension：最大 256（以 FlashAttention-2 为准）
attention mask：支持 boolean mask 或 is_causal=True，不支持任意的 float mask

若不满足这些条件，PyTorch 会自动回退到 memory-efficient attention 或 math backend。

11.5 实践建议

# 检查当前使用的是哪个 backend
import torch.backends.cuda

# 检查各 backend 的启用状态
print(f"Flash SDP enabled: {torch.backends.cuda.flash_sdp_enabled()}")
print(f"Mem efficient SDP enabled: {torch.backends.cuda.mem_efficient_sdp_enabled()}")
print(f"Math SDP enabled: {torch.backends.cuda.math_sdp_enabled()}")

# 全局禁用特定 backend
torch.backends.cuda.enable_flash_sdp(False)  # 禁用 FlashAttention
torch.backends.cuda.enable_mem_efficient_sdp(True)

11.6 直接使用 flash-attn 库

除了 PyTorch 原生的 SDPA 之外，也可以直接使用 Tri Dao 的 flash-attn 包。相比 PyTorch SDPA，它提供了更多功能（例如 sliding window attention、ALiBi、cross-attention 优化）：

# pip install flash-attn
from flash_attn import flash_attn_func

# 形状为 (batch, seqlen, nheads, headdim)
output = flash_attn_func(q, k, v, causal=True)

12. 总结与核心要点

FlashAttention 给出的核心教训是：算法的 FLOP 复杂度并非唯一决定性能的因素。在现代 GPU 上，内存访问模式才是主导实际执行时间的关键，IO-aware 的算法设计对实际性能起着决定性作用。

主要贡献可归纳如下：

IO-Aware 设计原则：利用 GPU 内存层级（HBM vs SRAM）非对称性的算法设计
Tiling + Online Softmax：以适配 SRAM 的块为单位运算，消除 $N \times N$ 矩阵在 HBM 中的具体化
Recomputation 策略：在 backward pass 中重新计算中间值，把内存从 $O(N^2)$ 降到 $O(N)$ ，同时还提升了速度
最优性证明：从 IO complexity 的角度证明下界，证实了算法的最优性
Exact Computation：在所有优化之上，仍保持不做近似的 exact attention

FlashAttention 是理论之美与实用效果兼具的罕见研究，已成为现代 LLM 训练与推理的核心基础设施。由于其已原生集成进 PyTorch，无需额外实现，仅调用 F.scaled_dot_product_attention 即可享受到它带来的收益。

References

Dao, T., Fu, D.Y., Ermon, S., Rudra, A., & Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. https://arxiv.org/abs/2205.14135
Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. https://arxiv.org/abs/2307.08691
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. NeurIPS 2024 Spotlight. https://arxiv.org/abs/2407.08608
Dao-AILab. flash-attention GitHub Repository. https://github.com/Dao-AILab/flash-attention
PyTorch Documentation. torch.nn.functional.scaled_dot_product_attention. https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
PyTorch Documentation. torch.nn.attention.sdpa_kernel. https://docs.pytorch.org/docs/stable/generated/torch.nn.attention.sdpa_kernel.html
PyTorch Blog. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. https://pytorch.org/blog/flashattention-3/
Milakov, M. & Gimelshein, N. (2018). Online Normalizer Calculation for Softmax. arXiv:1805.02867. https://arxiv.org/abs/1805.02867
NVIDIA. A100 Tensor Core GPU Architecture Whitepaper. https://images.nvidia.com/aem-dam/en-zz/Solutions/data-center/nvidia-ampere-architecture-whitepaper.pdf
NVIDIA. Hopper Architecture In-Depth. https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/