一个反直觉的事实：大模型推理的瓶颈不是计算，而是显存

如果你问一个外行：“大模型推理为什么这么慢？“他可能会说：“因为计算量太大了，需要算很多矩阵乘法。”

但真相是：大模型推理的瓶颈不是计算，而是显存带宽。 GPU的计算速度（FLOPS）远快于显存读写速度（Memory Bandwidth）。H100的计算速度是2000 TFLOPS（FP16），但显存带宽只有3.35 TB/s。计算速度是显存带宽的600倍。

这意味着，GPU花在"等待数据从显存加载到计算单元"上的时间，远多于"真正计算"的时间。FlashAttention解决的就是这个问题——减少显存读写，让计算单元不再"空转等待”。

Self-Attention的计算瓶颈

标准的Self-Attention计算流程是这样的：

  1. 从显存读取Q、K、V矩阵（3个N×d的矩阵）
  2. 计算 Q×K^T → 得到注意力分数矩阵 S（N×N的矩阵）
  3. 对S做Softmax
  4. 将Softmax结果 × V → 得到输出矩阵 O（N×d的矩阵）
  5. 将O写回显存

问题出在第2步：S矩阵的大小是N×N（N是序列长度）。 对于128K上下文，N=128000，S矩阵有160亿个元素，需要约32GB显存（FP16）。而且S矩阵需要反复读写显存——先写进去，再读出来做Softmax，再写进去，再读出来乘V。

每次做Self-Attention，都要读写一个巨大的S矩阵。 这就是显存带宽瓶颈的根源。

FlashAttention的核心创新：分块计算

FlashAttention的核心思想是：不要把整个S矩阵写回显存，而是分块计算、就地累加。

具体来说：

  1. 将Q分成小块（Block），每次只加载一个Q块到SRAM（GPU的片上缓存，非常快但很小）
  2. 将K、V也分成小块，逐个加载到SRAM
  3. 在SRAM中计算局部的注意力分数，做Softmax，乘V，累加到输出
  4. 处理完所有K、V块后，将最终输出写回显存

关键：S矩阵（N×N）永远不会被完整写回显存。 它只在SRAM中存在，用完就丢弃。这就省掉了最大的一笔显存读写开销。

效果：显存读写量从 O(N^2) 降低到 O(N)。 对于128K上下文，显存读写量降低约1000倍。

FlashAttention-3的进一步增强

2026年，FlashAttention已经发展到第3代（FlashAttention-3）。相比FlashAttention-2，FA3做了两个关键增强：

1. 异步计算。 FA3利用H100的新特性（TMA，Tensor Memory Accelerator），将数据加载和计算完全异步化。当计算单元在处理当前块时，显存控制器已经在加载下一个块。计算和显存读写不再"串行等待”，而是"并行重叠"。

2. FP8支持。 FA3原生支持FP8精度，可以在不损失质量的情况下，将显存占用再降低一半，速度再提升30%。

FA3在H100上的实测效果（Llama 4 70B，128K上下文）：

  • 推理速度：FA2的1.5倍，标准注意力的4倍
  • 显存占用：标准注意力的1/5

FlashAttention的局限性

FlashAttention不是"免费午餐"。 它有三个局限性：

1. 只适用于Self-Attention。 FlashAttention只优化了注意力机制这一部分。对于FFN层、Embedding层、LayerNorm层，FlashAttention无能为力。

2. 对硬件有要求。 FA3需要H100的TMA特性，在A100上只能使用FA2。FA2需要SM80+的GPU（A100、RTX 3090+）。

3. 分块计算有精度损失。 分块Softmax的数值精度略低于完整Softmax（误差通常在1e-5量级）。对于绝大多数应用，这个误差可以忽略。但对于一些对精度极度敏感的场景（如科学计算），可能需要关闭FlashAttention。

如何在你的代码中使用FlashAttention？

最简单的方法：使用支持FlashAttention的推理框架。

  • vLLM：默认启用FlashAttention
  • SGLang：默认启用FlashAttention
  • TensorRT-LLM：默认启用FlashAttention

如果你自己写推理代码：

# PyTorch 2.0+ 自动启用FlashAttention
torch.nn.functional.scaled_dot_product_attention(query, key, value)
# PyTorch会自动选择最优的注意力实现（FlashAttention > Memory Efficient Attention > Standard Attention）

结语：FlashAttention是"无声的革命"

FlashAttention是一项"无声的革命"——用户可能不知道它是什么，但享受了它带来的速度提升。 2026年，几乎所有主流推理框架都默认启用了FlashAttention。它让128K长上下文的推理从"不可能"变成了"可能"，从"巨慢"变成了"快"。

FlashAttention的启示：AI的瓶颈不是算法，而是硬件。 最聪明的优化，不是在算法层面"减少计算量"，而是在硬件层面"减少显存读写"。理解硬件，才能真正理解AI的性能。

数据来源：FlashAttention论文（Dao et al., 2022-2024），FlashAttention-3论文（2025），NVIDIA H100白皮书。