FlashAttention 通过分块计算和算子融合，将注意力内存复杂度从 $O(N^2)$ 降至 $O(N)$，实现 2-4 倍加速。

核心原理

内存层次优化

核心思想：将计算限制在 SRAM 中，减少 HBM 访问。

分块计算

• Q 分为 $T_r$ 块，K/V 分为 $T_c$ 块
• 每块大小 $B_r = B_c = 64$
• 局部注意力计算在 SRAM 中完成

在线 Softmax

单遍扫描更新全局统计量：

$$\begin{aligned} m_{\text{new}} &= \max(m_{\text{old}}, m_{ij}) \ l_{\text{new}} &= e^{m_{\text{old}} - m_{\text{new}}} \cdot l_{\text{old}} + e^{m_{ij} - m_{\text{new}}} \cdot l_{ij} \end{aligned}$$

版本演进

性能对比

使用方式

1
2
3
4
5
6

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B",
    use_flash_attention_2=True
)

或使用 vLLM 自动启用：

1
2

from vllm import LLM
llm = LLM(model="meta-llama/Llama-3-8B")

应用场景

• LLM 训练（LLaMA、ChatGLM）
• 长上下文处理（32k+ tokens）
• 高并发推理服务
• 多模态任务

资源链接

• 论文
• 官方实现