FlashAttention

FlashAttention 是一种用于计算 Transformer 模型中 attention 机制的优化算法。与标准 attention 相比，它速度更快、内存效率更高、可扩展性也更好。自论文 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness 提出以来，它已成为现代 LLM 在训练和推理中广泛采用的 attention backend。

FlashAttention 是一种 kernel 层面的优化。它的收益来自 attention 计算本身在 GPU 上的执行方式， 包括数据如何在 HBM 与 SRAM 之间移动、工作如何进行 tiling，以及操作如何融合进单个 kernel。 这正是本页放在 Kernel Optimization 章节中的原因。关于服务层的优化技术，请参见 Inference Optimization.

attention 为什么一开始就很慢

当 LLM 读取文本时，它需要查看每个 token，并将其与其他每个 token 进行比较，以理解它们之间的关系。 这被称为 attention.

标准 attention 机制有一个根本问题： 它是 memory-bound，而不是 compute-bound。要理解这一点，我们需要看看 attention 计算期间究竟发生了什么。

如果你想了解 HBM、SRAM、warp、tiling 等概念背后的底层 GPU 背景，可以参见 GPU architecture fundamentals.

标准 attention 计算的是：

$\text{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

朴素实现通常遵循以下步骤：

1. 计算 attention score：将 Q 与 K^T 相乘，得到一个 N×N 矩阵 （其中 N 是序列长度）
2. 应用 softmax：对 score 做归一化
3. 与 values 相乘：用 V 计算加权和

问题出在内存访问模式上。现代 GPU 具有：

• High Bandwidth Memory (HBM)：容量大，但速度较慢（1-2 TB/s 带宽）
• SRAM（片上内存）：容量小，但速度快（10-20 TB/s 带宽）

标准实现需要先将完整的 N×N attention 矩阵写入 HBM，然后再为下一步操作把它读回来。对于 4096 token 的序列长度，这个 attention 矩阵大约包含 1600 万个元素。由于存在多次读写， 算法大部分时间都耗在等待内存传输上，而不是做真正的计算。

随着序列长度增加：

• memory traffic 会主导运行时间
• GPU utilization 会下降
• 长上下文窗口会变得不切实际

例如，16K token 所需的内存是 1K token 的 256 倍。

FlashAttention 如何工作？

FlashAttention 通过减少 memory traffic 来加速 attention。其核心思想是永远不在 HBM 中 materialize 完整的 attention 矩阵。相反，它使用两个关键技术：

• Tiling 与 recomputation。FlashAttention 将计算拆分为能装入高速 SRAM 的 block（tile）：
  • 将 Q、K、V 的 tile 从 HBM 加载到高速 SRAM
  • 完全在 SRAM 中完成该 tile 的 attention 计算
  • 以增量方式更新输出，并丢弃中间结果
• Kernel fusion。FlashAttention 不再把操作拆成独立步骤（matmul → softmax → matmul），而是将其全部融合到单个 GPU kernel 中。这意味着：
  • 不需要将中间结果写入 HBM
  • 不需要单独发起多个 kernel launch（这本身有开销）
  • 所有操作都在高速 SRAM 中完成

FlashAttention 通过 tiling 避免在 HBM 上 materialize 巨大的 N×N attention 矩阵。图片来源

简单来说，FlashAttention 让 attention 计算更高效。它重新组织了工作方式，使 GPU 把更少时间花在等待内存上，把更多时间用于实际计算。

如果想从更宏观的角度了解 Triton、CUDA、compiler stack 以及 profiling tools 在这类工作中的定位，请参见 Kernel optimization tools.

FlashAttention 的收益

FlashAttention 在速度和可扩展性两方面都带来了显著提升：

• attention 速度提升 2–4 倍
• 由于不再存储 N×N attention 矩阵，内存占用大幅降低
• 让 LLM 能处理更长的上下文窗口（例如 128K token）
• 更高的吞吐量和更好的 GPU utilization
• 为聊天、代码生成、推理等任务带来更快的推理速度

目前，FlashAttention 已广泛用于：

• 训练框架（PyTorch、DeepSpeed）
• 推理引擎（vLLM、SGLang、Hugging Face TGI、TensorRT-LLM）
• 支持长上下文的模型架构

FlashAttention 版本对比

FlashAttention 主线目前已有 4 个主要版本。下表并排对比了该算法在不同版本中的演进方式。

版本	年份	关键改进	性能	说明
FlashAttention-1	2022	引入了具备 IO-awareness 的 tiled attention 算法。融合 softmax + matmul kernel。避免 materialize 完整 attention 矩阵	attention 速度提升 2–4 倍，内存占用最多可降低 10 倍	第一个版本；使长上下文在实践中可用；精确 attention（非近似）
FlashAttention-2	2023	改进 warp 间并行性与工作划分；减少非 matmul FLOPs	相比 FA-1 最多快 2 倍，尤其是在长序列上	驱动了许多长上下文 LLM；已广泛集成到推理/训练框架中
FlashAttention-3	2024	Tensor core 加速（FP8/BF16）；面向 Hopper GPU（如 H100）优化	相比 FA-2 最多快 2 倍，并在 H100 上达到 740 TFLOPS（75% util）；FP8 数值误差降低 2.6 倍	利用 Hopper 的异步执行与 warp specialization。许多框架仍然会先从 FA-2 升级
FlashAttention-4	2026	完全异步的 MMA、更大的 tile、软件模拟指数、条件式 softmax rescaling、tensor memory，以及 2-CTA MMA	在 B200 BF16 基准上，相比 cuDNN 9.13 最多快 1.3 倍、相比 Triton 最多快 2.7 倍，并达到最高 1613 TFLOPS/s（71% utilization）	使用 CuTeDSL 编写。官方实现通过 flash-attn-4 暴露，目标为 H100、B200 等 Hopper 与 Blackwell GPU

FlashAttention-4 是专门针对 NVIDIA Blackwell 架构调优的。其关键洞见是硬件扩展具有不对称性。 Tensor core（负责 QKᵀ 和 PV 这类大型矩阵乘法）在 Blackwell 上变得更快了，但其他关键资源的扩展幅度并没有那么大：

• Shared memory 带宽
• softmax 中用于指数运算的 special function units（SFU）
• register pressure 与调度开销

因此，在 B200 上，瓶颈发生了转移。详情可参见 FlashAttention-4 paper.

如何使用 FlashAttention

最容易的入门方式是使用官方包：

pip install flash-attn --no-build-isolation

较新的 PyTorch 版本在支持时会通过 scaled_dot_product_attention 自动分发到 FlashAttention。 更多信息请阅读 API 参考文档。

许多推理框架已经集成了 FlashAttention，包括 vLLM 和 SGLang， 但由于各自的发布节奏不同，它们所集成的版本可能并不一致。

对于 FlashAttention-4，官方仓库说明了一个单独的 CuTeDSL 包：

pip install flash-attn-4

在 CUDA 13 上，仓库推荐：

pip install "flash-attn-4[cu13]"

FA4 API 通过 flash_attn.cute 暴露：

from flash_attn.cute import flash_attn_func

out = flash_attn_func(q, k, v, causal=True)

在生产环境中依赖它之前，请先查看最新的 FlashAttention 仓库，因为 FA4 包及其框架集成进展很快。

延伸资源

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
• FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling
• FlashAttention 官方仓库