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)理论上限
因果推理4096215 TFLOPS350 TFLOPS710 TFLOPS989
因果推理8192235 TFLOPS420 TFLOPS780 TFLOPS989
因果推理16384245 TFLOPS460 TFLOPS820 TFLOPS989
全注意力8192280 TFLOPS520 TFLOPS890 TFLOPS989

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

4.3 对比 cuDNN Attention

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

指标cuDNN 9.0FA3差异
FP16 推理380 TFLOPS460 TFLOPS+21%
FP8 推理650 TFLOPS820 TFLOPS+26%
长序列 (32K)310 TFLOPS440 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%
FA12022分块 + 在线 softmax35-50%
FA22023减少非 GEMM + 并行优化50-73%
FA3 (FP16)2024异步执行 + TMA47-72%
FA3 (FP8)2024FP8 Tensor Core72-90%
FA3.5 (FP4)2026FP4 + DSMEM预计 80-95%

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

加入讨论

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

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