AI Infra：FlashAttention，具有I/O感知的快速且内存高效的精确注意力机制
一、FlashAttention 简介
FlashAttention 是由 Stanford Hazy Research 团队提出的高效注意力机制实现，核心目标是在不牺牲准确率的前提下最大化注意力计算效率，尤其适用于 GPU 上的训练与推理。
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 表现更稳 |