Flash Attention 3：让 GPU 跑满的注意力计算

Flash Attention 系列是近年来大模型工程领域最重要的优化之一。从 Flash Attention 1 到 3，每代都在逼近 GPU 硬件的理论极限。2026 年，Flash Attention 3 已经成为所有主流大模型推理和训练的标配。本文将深入解析其原理。

一、问题：标准注意力的内存瓶颈

1.1 GPU 内存层次

现代 GPU（如 H100）有复杂的内存层次：

┌─────────────────────────────────────────────┐
│              GPU 内存层次                    │
├─────────────────────────────────────────────┤
│                                             │
│  ┌─────────┐  延迟: ~20 cycles             │
│  │Register │  带宽: ~30 TB/s               │
│  │(SRAM)   │  容量: 256 KB/SM              │
│  └─────────┘                                │
│       ↑                                     │
│  ┌─────────┐  延迟: ~200 cycles            │
│  │L2 Cache │  带宽: ~12 TB/s               │
│  │         │  容量: 50 MB                  │
│  └─────────┘                                │
│       ↑                                     │
│  ┌─────────┐  延迟: ~400+ cycles           │
│  │  HBM    │  带宽: ~3.35 TB/s (H100)      │
│  │(显存)   │  容量: 80 GB                  │
│  └─────────┘                                │
│                                             │
└─────────────────────────────────────────────┘

关键洞察：HBM 带宽只有 SRAM 的 1/9，但标准注意力几乎完全在 HBM 上操作。

1.2 标准注意力的内存访问模式

# 标准注意力实现
def attention_standard(Q, K, V):
    # Q, K, V: [n, d] in HBM
    S = Q @ K.T          # [n, n] — 中间矩阵在 HBM，需要 O(n²) 读写
    P = softmax(S)       # [n, n] — 中间矩阵在 HBM
    O = P @ V            # [n, d] — 输出在 HBM
    return O

HBM 读写量：

  • 读 Q, K, V：$3 \times n \times d$
  • 写 S：$n^2$
  • 读 S, 写 P：$2 \times n^2$
  • 读 P, 写 O：$n^2 + n \times d$

总计：$\Theta(n^2 + n \cdot d)$，其中 $n^2$ 项是瓶颈。

1.3 核心矛盾

注意力矩阵 $S \in \mathbb{R}^{n \times n}$ 太大，无法放入 SRAM（通常只有几 MB），所以必须写入 HBM。但 HBM 带宽有限，成为计算瓶颈。

Flash Attention 的核心思想：不把完整的 $n \times n$ 矩阵写入 HBM，而是在 SRAM 中分块计算。

二、Flash Attention 1：分块计算

2.1 算法核心

将 Q, K, V 分块加载到 SRAM，在 SRAM 中完成注意力计算：

┌────────────────────────────────────────────────┐
│          Flash Attention 分块策略               │
├────────────────────────────────────────────────┤
│                                                │
│  Q: ┌──┬──┬──┬──┐                             │
│     Q1 Q2 Q3 Q4   (每块 Br 行)                 │
│                                                │
│  K: ┌──┬──┬──┬──┐                             │
│     K1 K2 K3 K4   (每块 Bc 行)                 │
│                                                │
│  V: ┌──┬──┬──┬──┐                             │
│     V1 V2 V3 V4                               │
│                                                │
│  外层循环: for j = 1..4:                       │
│    加载 Kj, Vj 到 SRAM                         │
│    内层循环: for i = 1..4:                     │
│      加载 Qi 到 SRAM                           │
│      在 SRAM 中计算 Sij = Qi @ Kj.T             │
│      在 SRAM 中计算 softmax(Sij) @ Vj           │
│      累加结果到 Oi                              │
│      (只写最终 Oi 到 HBM)                      │
│                                                │
└────────────────────────────────────────────────┘

2.2 分块 Softmax 的数学

标准 softmax：$\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}$

分块计算需要处理跨块的全局归一化。使用在线 softmax算法：

对于块 $j$，维护运行统计量 $m$（最大值）和 $l$（求和）：

$$m_{new} = \max(m_{old}, m_j)$$ $$l_{new} = l_{old} \cdot e^{m_{old} - m_{new}} + l_j \cdot e^{m_j - m_{new}}$$ $$O_{new} = O_{old} \cdot \frac{l_{old} \cdot e^{m_{old} - m_{new}}}{l_{new}} + P_j \cdot V_j \cdot \frac{e^{m_j - m_{new}}}{l_{new}}$$

这个数学技巧让分块 softmax 与标准 softmax 在数值上等价。

2.3 IO 复杂度

  • 标准注意力：$\Theta(n^2 + nd)$ HBM 读写
  • Flash Attention：$\Theta(n^2 d / M)$ HBM 读写，其中 $M$ 是 SRAM 大小

当 $M \geq d$ 时（总是成立），Flash Attention 的 IO 量远小于标准方法。

三、Flash Attention 2：减少非矩阵乘法

3.1 优化要点

FA1 的瓶颈在于非 GEMM 操作（softmax、缩放、统计量更新）占比过高。FA2 重新排列了计算顺序：

  1. 减少非 GEMM FLOPs：将缩放因子合并到 Q 的初始化中
  2. 优化并行度：在序列维度和外层循环上并行
  3. 减少同步：每个线程块独立处理一个 Q 块

3.2 性能提升

在 A100 上，FA2 达到了理论 FLOPS 的 50-73%（标准实现仅 25-35%）。

四、Flash Attention 3：硬件感知优化

4.1 核心创新

Flash Attention 3（2024-2025）针对 Hopper 架构（H100）进行了深度优化，三大创新：

4.1.1 异步执行（Warp Specialization）

H100 支持 TMA（Tensor Memory Accelerator），可以异步在 HBM 和 SRAM 之间传输数据。FA3 利用这一点实现了生产者-消费者模式：

┌─────────────────────────────────────────────┐
│         Warp Specialization (4 warps)       │
├─────────────────────────────────────────────┤
│                                             │
│  Warp 0 (生产者): TMA 加载 K, V → SRAM      │
│  Warp 1-3 (消费者): Tensor Core 计算 QKV    │
│                                             │
│  时间线:                                    │
│  ┌───────────────────────────────────────┐  │
│  │ Load K1│Load K2│Load K3│Load K4│     │  │
│  │────────┼───────┼───────┼───────┼     │  │
│  │Calc Q1 │Calc Q2│Calc Q3│Calc Q4│     │  │
│  │  K1    │  K2   │  K3   │  K4   │     │  │
│  └───────────────────────────────────────┘  │
│  数据加载和计算完全重叠!                      │
│                                             │
└─────────────────────────────────────────────┘

4.1.2 FP8 支持

H100 的 FP8 Tensor Core 吞吐量是 FP16 的 2 倍（1979 TFLOPS vs 989 TFLOPS）。FA3 支持两种 FP8 格式：

  • E4M3（4位指数，3位尾数）：精度更高，用于前向传播
  • E5M2（5位指数，2位尾数）：动态范围更大，用于反向传播

FP8 注意力的关键挑战是 softmax 的数值稳定性。FA3 使用两阶段 softmax：

  1. 先用 FP16 计算最大值和求和（保证稳定性）
  2. 再用 FP8 执行矩阵乘法（利用高吞吐）

4.1.3 低精度 KV Cache

FA3 原生支持 INT8 和 FP8 KV Cache，在不修改推理逻辑的前提下减少内存占用。

4.2 性能基准

在 H100 80GB SXM5 上的测试结果：

配置 序列长度 FA2 (FP16) FA3 (FP16) FA3 (FP8) 理论上限
因果推理 4096 215 TFLOPS 350 TFLOPS 710 TFLOPS 989
因果推理 8192 235 TFLOPS 420 TFLOPS 780 TFLOPS 989
因果推理 16384 245 TFLOPS 460 TFLOPS 820 TFLOPS 989
全注意力 8192 280 TFLOPS 520 TFLOPS 890 TFLOPS 989

FA3 FP8 在长序列上达到了理论峰值 90% 的利用率。

4.3 对比 cuDNN Attention

NVIDIA 官方的 cuDNN Attention 在 H100 上的表现：

指标 cuDNN 9.0 FA3 差异
FP16 推理 380 TFLOPS 460 TFLOPS +21%
FP8 推理 650 TFLOPS 820 TFLOPS +26%
长序列 (32K) 310 TFLOPS 440 TFLOPS +42%
内存占用 基准 -15% 更优

FA3 在所有场景下都优于 cuDNN，尤其是长序列场景。

五、Flash Attention 3 的工程实现

5.1 关键参数调优

# FA3 的关键参数
block_size_q = 128    # Q 块大小
block_size_k = 128    # K 块大小
block_size_v = 128    # V 块大小
num_warps = 4         # warp 数量
num_stages = 3        # 流水线阶段数

# 参数选择经验:
# - 短序列 (n < 4K): 大块 (256), 少 warp (2)
# - 中序列 (4K-32K): 中块 (128), 4 warp
# - 长序列 (>32K): 小块 (64), 多 warp (8), 多阶段 (4)

5.2 变长序列处理

实际场景中 batch 内不同请求的长度不同。FA3 使用 varlen 模式：

# 变长序列的 FA3 调用
output = flash_attn_3_func(
    q, k, v,
    cu_seqlens_q=cu_seqlens,  # 累积长度 [0, len1, len1+len2, ...]
    cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_len,
    max_seqlen_k=max_len,
    causal=True,
    softmax_scale=1.0 / sqrt(d_k),
)

这避免了 padding 造成的计算浪费。

5.3 与 PagedAttention 集成

FA3 可以与 vLLM 的 PagedAttention 集成，实现分页 KV Cache + Flash 计算：

┌────────────────────────────────────────────┐
│       FA3 + PagedAttention 流程            │
├────────────────────────────────────────────┤
│                                            │
│  1. 从页表中收集 KV Cache 页面             │
│  2. 连续化到临时缓冲区                      │
│  3. 在连续缓冲区上执行 FA3                  │
│  4. 写回结果                               │
│                                            │
│  性能影响: 约 5-10% 额外开销               │
│                                            │
└────────────────────────────────────────────┘

六、2026 年展望：Flash Attention 在 Blackwell 上

6.1 Blackwell 架构的新特性

NVIDIA B200 GPU（Blackwell 架构）带来了新的硬件特性：

  • FP4 Tensor Core：支持 4-bit 浮点计算
  • 2nd Gen TMA：更高效的异步传输
  • Distributed Shared Memory：多个 SM 可以共享 SRAM

6.2 Flash Attention 3.5（预览）

针对 Blackwell 的优化版本已经在开发中：

优化方向 预期提升
FP4 注意力 2x 吞吐 vs FP8
DSMEM 跨 SM 通信 减少 30% 同步开销
2nd Gen TMA 减少 20% 数据加载延迟
混合精度策略 FP4 计算 + FP8 累积

预计 FA3.5 在 B200 上可以达到 3.5 PFLOPS 的注意力计算吞吐。

七、总结

Flash Attention 的发展历程是一个完美的"硬件-算法协同设计"案例：

版本 年份 核心创新 GPU 利用率
标准实现 2017 无 25-35%
FA1 2022 分块 + 在线 softmax 35-50%
FA2 2023 减少非 GEMM + 并行优化 50-73%
FA3 (FP16) 2024 异步执行 + TMA 47-72%
FA3 (FP8) 2024 FP8 Tensor Core 72-90%
FA3.5 (FP4) 2026 FP4 + DSMEM 预计 80-95%

Flash Attention 证明了一个重要原则：算法设计必须理解硬件。在 GPU 上，减少内存访问往往比减少计算量更重要。这一教训适用于所有大规模 AI 系统的设计。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。