AI Infra:FlashAttention,具有I/O感知的快速且内存高效的精确注意力机制

新闻

作者：微信文章



一、FlashAttention 简介

FlashAttention 是由 Stanford Hazy Research 团队提出的高效注意力机制实现，核心目标是在不牺牲准确率的前提下最大化注意力计算效率，尤其适用于 GPU 上的训练与推理。

https://github.com/Dao-AILab/flash-attention

Paper: https://tridao.me/publications/flash3/flash3.pdf

1.1 核心原理

• 将传统的多次读取/写入的注意力计算过程，融合为一次 GPU kernel 调用• 使用 tile-based 的块级处理方式，在寄存器中缓存中间变量，大大减少显存访问和带宽瓶颈• 避免了中间激活的 O(n²) 存储，特别适合长序列

---

1.2 版本说明

版本	特点
FlashAttention v1	用于训练/推理，速度提升 2~4 倍；兼容性好
FlashAttention v2	全 kernel 实现，支持 dropout/causal，进一步减少显存读写；支持更长序列
FlashAttention v3（2024）	针对 NVIDIA Hopper 优化，新增对 Multi-Query Attention 的支持，支持推理时 KV cache，显著提升推理吞吐（约 1.3–1.5×）

---

1.3 表现

GPU	Speedup (vs PyTorch)	Memory Saving
A100	2–4×	长序列下节省 10–20× 内存
RTX 3090	2.5–4.5×	同上
T4	2.5–4.5×（仅 fwd）	比较节省

二、与 FireAttention 的对比

2.1 什么是 FireAttention

FireAttention 是由 Fireworks AI 自研的高性能 CUDA 注意力内核，专为加速多查询注意力（Multi‑Query Attention）模型（如 Mixtral 等 MoE 架构）而打造。

它集成到 Fireworks 的 LLM 推理堆栈中，支持 FP16、FP8、FP4 等多种精度格式，显著提升推理速度与吞吐能力，同时保持模型质量

客户有 poe.com

https://fireworks.ai/

---

2.2 对比表

比较维度	FireAttention	FlashAttention
开发者	Fireworks AI	Stanford Hazy Research 团队
初始发布时间	2024 初	2022（v1），v3 发布于 2024
设计目标	超高吞吐 & 多精度推理优化，专注于大模型服务场景	通用型高效注意力内核，适用于训练和推理
优化方向	专为推理优化（特别是 MoE/MQA），支持 FP8/FP4	训练与推理通用优化，面向长序列 & 更大 batch
性能指标（推理）	最多 4× vLLM，FP8/FP4 显著提速（可达 250+ token/s）	约 1.3–1.5× 吞吐提升（v3）；更节省内存
精度支持	FP16 / FP8 / FP4（V4 支持 Blackwell）	FP16 / BF16，v3 有 FP8 支持（实验中）
MQA / MoE 支持	强（如 Mixtral）	v3 开始支持，尚未深度集成于推理堆栈
KV Cache 优化	高度优化，支持 token streaming 和动态 batch	v3 开始支持，仍以训练为主
芯片兼容性	NVIDIA H100 / B200 / AMD MI300	NVIDIA A100 / H100 等
应用平台/产品化	已应用于 Fireworks 全平台，服务 Quora Poe 等	多用于 Hugging Face、OpenAI 微调、训练等

---

2.3 适合场景

使用场景	推荐
通用 LLM 训练 + 推理平台	FlashAttention v2/v3：兼容好，长序列友好
高速云端 LLM 推理服务（MoE / 多租户）	FireAttention：推理吞吐最佳，支持低比特精度（FP4/FP8）
构建 Chatbot、RAG 工程化平台	FireAttention 更适合
长文本/预训练大模型（如 LLaMA）	FlashAttention 表现更稳