使用 FlashInfer 加速 LLM 服务中的自注意力机制

LLM（大型语言模型）服务迅速成为一项重要的工作负载。Transformer 中的操作符——即 GEMM、自注意力、GEMV 和逐元素计算——的效率对于 LLM 服务的整体性能至关重要。虽然优化工作已广泛针对 GEMM 和 GEMV，但在 LLM 服务中缺乏对自注意力的性能研究。在这篇博客文章中，我们将自注意力分解为三个阶段：预填充（prefill）、解码（decode）和追加（append）；分析在这三个阶段中，自注意力在单请求和批量处理场景下的性能瓶颈；并提出解决这些挑战的方案。这些想法已整合到 FlashInfer 中，这是一个在 Apache 2.0 许可下发布用于加速 LLM 服务的开源库。

FlashInfer 自2023年夏天以来由来自华盛顿大学、卡内基梅隆大学和 OctoAI 的研究人员开发。FlashInfer 提供 PyTorch API 以进行快速原型设计，以及一个无依赖的、仅包含头文件的 C++ API，用于与 LLM 服务系统集成。与现有库相比，FlashInfer 具有几个独特的优势

1. 全面的注意力核函数：FlashInfer 实现了注意力核函数，涵盖了 LLM 服务的所有常见用例，并具有最先进的性能，包括用于预填充、解码和追加核函数的单请求和批量处理版本，支持各种格式的 KV-Cache（填充张量、不规则张量和页表）。
2. 优化的共享前缀批量解码：FlashInfer 通过级联（cascading）增强了共享前缀批量解码性能，与基准 vLLM PageAttention 实现相比，实现了高达31倍的加速（对于32768令牌的长提示和256的大批量大小），请参阅另一篇博客文章了解更多详情。
3. 加速压缩/量化 KV-Cache 的注意力机制 现代 LLM 通常部署带有量化/压缩 KV-Cache 以减少内存流量。FlashInfer 通过优化分组查询注意力（Grouped-Query Attention）、融合 RoPE 注意力（Fused-RoPE Attention）和量化注意力（Quantized Attention）的性能来加速这些场景。值得注意的是，与 vLLM 实现相比，FlashInfer 在 A100 和 H100 上实现了高达 2-3 倍的分组查询注意力加速。

FlashInfer 已被 LLM 服务系统采用，例如 MLC-LLM（用于其 CUDA 后端）、Punica 和 sglang。我们欢迎社区更广泛地采用和贡献。请加入我们的讨论论坛或创建议题留下您的反馈和建议。

LLM 服务中的注意力机制

LLM 服务中有三个通用阶段：预填充（prefill）、解码（decode）和追加（append）。在预填充阶段，注意力计算发生在 KV-Cache 和所有查询之间。在解码阶段，模型一次生成一个令牌，只计算 KV-Cache 和单个查询之间的注意力。在追加阶段，注意力在 KV-Cache 和追加的令牌的查询之间计算。追加注意力在推测性解码（speculative decoding）中也很有用：草稿模型建议一个令牌序列，而更大的模型决定是否接受这些建议。在注意力阶段，提出的令牌被添加到 KV-Cache 中，大模型计算 KV-Cache 和提出的令牌之间的注意力。

影响注意力计算效率的关键因素是查询长度 ($l_q$)，它决定了操作是计算密集型（compute-bound）还是 IO 密集型（IO-bound）。注意力计算的操作强度（每字节内存流量的操作数）表示为 $O\left(\frac{1}{1/l_q + 1/l_{kv}} \right)$，其中 $l_{kv}$ 表示 KV-Cache 的长度。在解码阶段，其中 $l_q$ 始终为1，操作强度接近 $O(1)$，使操作符完全是 IO 密集型。在追加/预填充阶段，注意力操作强度大约为 $O(l_q)$，当 $l_q$ 较大时导致计算密集型场景。

该图说明了预填充、追加和解码阶段的注意力计算过程

图1：解码注意力一次填充注意力图的一行，预填充注意力填充整个注意力图（在因果掩码下），追加注意力填充梯形区域。

下图显示了注意力计算三个阶段的屋顶线模型。解码注意力性能始终低于峰值带宽上限（受 GPU 峰值内存带宽限制），因此是 IO 密集型。预填充注意力具有高操作强度，并且低于峰值计算性能上限（受峰值浮点性能限制）。追加注意力在查询长度较小时是 IO 密集型，在查询长度较大时是计算密集型。

图2. LLM 服务中注意力操作符的屋顶线模型，数据来自 A100 PCIe 80GB。

单请求和批量处理

有两种常见的方式来提供 LLM 模型服务：批量处理和单请求。批量处理将多个用户请求组合在一起并并行处理以提高吞吐量，然而，注意力核函数的操作强度与批量大小无关¹，批量解码注意力仍然具有 $O(1)$ 的操作强度。

FlashInfer 概览

FlashAttention 提出通过将 在线 softmax 技巧推广到自注意力来将多头注意力融合到一个核函数中，从而避免在 GPU 全局内存上实例化注意力矩阵的开销。FlashAttention2 通过采用更合理的平铺策略和减少非张量操作的数量来进一步提高性能，以减轻 A100/H100 非张量核心性能较低的问题。vLLM 提出了 PageAttention，其中 KV-Cache 被组织成一个页表，以减轻 LLM 服务中的内存碎片问题。

FlashInfer 实现了 FlashAttention 的单请求和批量版本，用于所有三个阶段：预填充、追加和解码，支持多功能 KV-Cache 格式（例如，不规则张量、页表）。对于单解码/预填充和批量解码核函数，FlashInfer 实现了单请求解码/预填充和批量解码核函数的最先进性能。此外，FlashInfer 实现了用于分页 KV-Cache 的预填充/追加核函数，这是现有库中前所未有的，并且可用于在推测性解码设置中服务模型。

最近的许多工作提出了 KV-Cache 压缩技术以减少内存流量。鉴于此，FlashInfer 优化了分组查询注意力、融合 RoPE 注意力和量化注意力的核函数，以实现高效的压缩 KV-Cache 服务

• 分组查询注意力：分组查询注意力为键和值使用较少数量的头，从而节省内存流量。分组查询注意力的操作强度从 $O(1)$ 增长到 $O\left(\frac{H_{qo}}{H_{kv}}\right)$，其中 $H_{qo}$ 是查询头的数量，$H_{kv}$ 是键和值头的数量。A100/H100 等 GPU 具有较低的非张量核心性能，因此传统的分组查询注意力实现是计算密集型的。FlashInfer 建议在 GQA 中使用预填充核函数（利用张量核心）进行解码注意力，与 vLLM 实现相比，实现了高达 2-3 倍的加速。
• 融合 RoPE 注意力：RoPE（旋转位置嵌入）已成为 Transformer 的标准组件，大多数现有的服务系统将 post-RoPE 键（应用旋转嵌入后的键）存储在 KV-Cache 中。然而，一些最近的工作（例如StreamingLLM）会修剪 KV-Cache 中的令牌，修剪后令牌的位置会发生变化，因此 KV-Cache 中的 post-RoPE 键变得毫无意义。在这种情况下，FlashInfer 建议在 KV-Cache 中保存 pre-RoPE 键，并将 RoPE 融合到注意力核函数中。在各种平台和设置上的实验表明，FlashInfer 的 Fused-RoPE Attention 核函数可以动态应用 RoPE，开销可忽略不计。
• 量化注意力：压缩 KV-Cache 的另一种方法是通过修剪，FlexGen 和 Atom 表明可以将 KV-Cache 修剪到4位，而精度损失可忽略不计。FlashInfer 实现了低精度注意力核函数，因此我们可以实现接近线性加速比的压缩比（4位约4倍，8位约2倍）。

一些最近的工作，例如 LightLLM 和 sglang，使用了一种特殊的 PageAttention 形式，其中页面大小等于一，以便在复杂的服务场景（例如结构化生成）中轻松管理 KV-Cache。FlashInfer 通过在 GPU 共享内存中预取页面索引来优化 PageAttention 核函数，从而使核函数性能不受页面大小影响。

在接下来的部分中，我们将深入研究 FlashInfer 实现的详细优化和基准测试结果。

基准测试设置

硬件

我们对4种不同的 GPU 进行了基准测试：H100 SXM 80GB、A100 PCIe 80GB、RTX 6000 Ada 和 RTX 4090，前两个是广泛使用的数据中心 GPU，分别基于 Hopper 和 Ampere 架构，后两个是工作站和游戏 GPU，基于 Ada Lovelace 架构，价格更实惠，规格列在下表中

	H100 SXM 80GB	A100 PCIe 80GB	RTX Ada 6000	RTX 4090
GPU 内存 (GB)	80	80	48	24
微架构	Hopper (sm_90)	Ampere (sm_80)	Ada Lovelace (sm_89)	Ada Lovelace (sm_89)
内存带宽 (GB/s)	3,352	1,935	960	1,008
SM 数量	132	108	142	128
峰值张量核心性能 (TFLops/s)	989	312	366	165 (f32累加) 330 (f16累加)
峰值 (非张量核心) FP32 性能 (TFLops/s)	67	20	90	80
最大共享内存 (KB/SM)	228	164	100	100
L2 缓存 (KB)	51200	40960	98304	73728

H100 SXM 80GB 使用 HBM3，A100 PCIe 80GB 使用 HBM2e，两者都比使用 GDDR6X 的 RTX Ada 6000 和 RTX 4090 具有更大的内存带宽。RTX Ada 6000 和 RTX 4090 具有比 A100（20 TFLops/s）大得多的非张量核心峰值性能（分别为90和80 TFLops/s）。后三款 GPU 在 f16 累加方面具有相似的峰值张量核心（fp16 输入，无稀疏性）性能，RTX 4090 的张量核心在 fp16 累加方面的吞吐量是 fp32 累加的两倍，而其他 GPU 在 fp16 和 fp32 累加方面具有相同的吞吐量。

以下是四款 GPU 的张量核心和 CUDA 核心的屋顶线曲线

图3：4款 GPU 的设备屋顶线，张量核心性能和 CUDA 核心性能分别显示。

拐点由峰值浮点性能与内存带宽的比率决定。

软件

进行比较的基准是：FlashAttention 2.4.2（已整合 FlashAttention 2 和 Flash-Decoding）和 vLLM v0.2.6（实现了 PageAttention 1&2）。对于 vLLM，我们使用 pip 预构建的 wheel，我们从源代码构建 FlashAttention 和 FlashInfer，使用 CUDA 12.3.1 版本中的 NVCC 编译器。核函数分析使用 nvbench 库完成，我们取“冷”GPU 时间，即在每次核函数启动前刷新 L2 缓存。

指标

我们报告预填充和追加注意力核函数达到的 TFLops/s，以及解码和追加注意力核函数的 GPU 内存带宽利用率（通过 $\frac{\textrm{核函数读取的字节数}}{\textrm{核函数延迟}} / \textrm{硬件 GPU 内存带宽}$ 计算）。

预填充核函数

对于预填充（多查询）注意力，我们重新实现了 FlashAttention 2 算法的纯 CUDA 版本，并进行了一些额外的优化。标准的 FlashAttention 实现使用 fp16 输入和 fp32 累加器的张量核心，然而，RTX 4090 GPU 在 fp32 累加器下的张量核心性能较低，我们观察到注意力计算中的 $\frac{\mathbf{q}\cdot \mathbf{k}^{T}}{\sqrt(d)}$ 阶段具有较小的取值范围，可以使用 fp16 累加（因为头维度总是很小：例如128），FlashInfer 提供了一个 allow_fp16_qk_reduction 选项来允许这种优化（但仍然对 $\mathbf{score} \cdot \mathbf{v}$ 使用 fp32 累加），这种优化可以在 RTX 4090 上带来50%的加速。下面是 FlashInfer 0.0.1 和 FlashAttention 2.4.2 在不同 GPU 上的性能比较

图4：单请求预填充核函数性能，使用 Llama2-7B 设置：num_kv_heads=num_qo_heads=32，head_dim=128。序列长度从32到65535不等。

在 f32 累加设置中，FlashInfer 的预填充核函数实现在所有4个 GPU 上均实现最佳性能。allow_fp16_qk_reduction 选项可以进一步提高性能，特别是对于 RTX 4090。

追加和解码优化

追加和解码注意力往往具有比查询长度更大的 KV 长度，这可能会在批量大小较小时限制 GPU 中的 SM（流多处理器）利用率，FlashInfer 提出在 GEMM 优化中使用 Split-K 技巧，通过在序列维度上拆分 KV-Cache 来增加并行度。另一项工作 Flash-Decoding 也探索了这一想法，您可以查看他们精彩的博客文章获取可视化和解释。下面是 FlashInfer 0.0.1 和 FlashAttention 2.4.2 在不同 GPU 上的解码注意力性能比较

图5：单请求解码核函数性能，使用 Llama2-7B 设置：num_kv_heads=num_qo_heads=32，head_dim=128。序列长度从32到65536不等。

FlashInfer 在所有4个 GPU 上均实现最佳性能，并且对于长序列，GPU 带宽利用率接近100%。一个有趣的发现是，split-KV 并不会提高 RTX Ada 6000 和 RTX 4090 等 GPU 的性能，因为它们具有相对较小的内存带宽和更强的 CUDA 核心性能（解码注意力具有低操作强度，并且我们在非 GQA 设置中使用 CUDA 核心）。与 SM 本地计算单元不同，GPU 上的全局内存流量是共享的，因此如果操作符不是计算密集型的，使用 32 个（Llama2-7B 设置中的头数）中的 108 个 SM 仍然可以充分利用内存带宽。A100 GPU 具有较低的 CUDA 核心性能（20 TFLops/s），使用 108 个 SM 中的 32 个（5.9 TFLops/s）将使核函数成为计算密集型（除了乘法和加法之外，在注意力计算中还有耗时的计算，例如 exp），在这种情况下 split-KV 是有帮助的。

对于批量解码注意力，FlashInfer 实现了带有预取页面索引等优化的 PageAttention，下面是 FlashInfer PageAttention 核函数和 vLLM PageAttention 核函数的性能比较

图6：批量解码核函数性能，使用 Llama2-7B 设置：num_kv_heads=num_qo_heads=32，head_dim=128，batch_size=[1,16,64]。对于 batch_size = 1，序列长度从32到65536不等，对于 batch_size = 16，从32到4096不等，对于 batch_size = 64，从32到1024不等。

在不同的批量大小和序列长度下，FlashInfer PageAttention 核函数相对于 vLLM 0.2.6 的实现具有持续的加速。

我们还对追加注意力核函数进行了基准测试（追加注意力与预填充注意力共享同一组 API，区别在于追加注意力中的查询长度小于键/值长度）

图7：追加注意力核函数性能，使用 Llama2-7B 设置，num_kv_heads=num_qo_heads=32，head_dim=128。追加长度设置为128或256，KV 序列长度从32到65536不等。

FlashInfer 在所有4个 GPU 上仍然实现了最佳性能，无论使用 fp16 还是 fp32 qk 累加器。Split-KV 显著提高了追加核函数的性能，对于追加长度为128和256都是如此，因为操作符的操作强度变大，使用 32/100+ SM 不再提供足够的计算单元，从而使核函数成为计算密集型。请注意，RTX 4090 的张量核心 fp32 累加器屋顶线的拐点是 163（165 TFLops/s / 1008 GB/s），当查询长度（大约等于操作强度）达到256时，核函数将成为计算密集型，使用 allow_fp16_qk_reduction 可以缓解这个问题。

FlashInfer 还实现了批量追加注意力核函数，其中键/值存储在页表中，这可以加速 LLM 服务中的推测性解码，我们将在另一篇博客文章中讨论。

分组查询注意力

分组查询注意力使用的键/值头数少于查询/输出头数，使得操作强度高于普通多头注意力。FlashInfer 建议在 GQA 中使用预填充（多查询）注意力核函数，该核函数利用张量核心进行解码注意力。下面是 FlashInfer（w/ CUDA 核心和 w/ 张量核心）和 FlashAttention 2.4.2 在 A100 和 H100 上的 GQA 核函数性能比较

图8：单请求 GQA 解码性能，使用 llama2-70b 设置：tp=2，num_kv_heads=4，num_qo_heads=32，head_dim=128。序列长度从32到65536不等。

对于单请求 GQA 解码注意力，FlashInfer（张量核心）在 A100 和 H100 上均实现比 FlashAttention 2.4.2 更好的性能，而 FlashInfer（CUDA 核心）由于有限的 CUDA 核心性能只能实现 40%+ 的带宽利用率。

图9：批量 GQA 解码性能，使用 llama2-70b 设置：tp=2，num_kv_heads=4，num_qo_heads=32，head_dim=128。batch_size 设置为64，每个请求的序列长度从32到8192不等。

对于批量 GQA 解码注意力，当 batch_size=64 时，带有张量核心的 FlashInfer 比 vLLM PagaAttention 快3倍。

融合 RoPE 注意力

KV-Cache 压缩技术（例如 H2O 和 Streaming-LLM）通过删除令牌来修剪 KV-Cache，KV-Cache 中令牌的原始相对位置会被污染，将 post-RoPE 键存储在 KV-Cache 中变得毫无意义。FlashInfer 实现了高性能的 Fused-RoPE 注意力核函数，可以动态应用 RoPE，下面是 FlashInfer 解码注意力有无 RoPE 的性能比较

图10：融合 RoPE 注意力性能，使用 Llama2-7B 设置：um_kv_heads=num_qo_heads=32，head_dim=128。序列长度从32到65536不等。

RoPE 在所有4个 GPU 上的开销可以忽略不计，特别是对于具有强大 CUDA 核心性能的 RTX 6000 Ada 和 RTX 4090 GPU（RoPE 需要 sin/cos 计算，不能通过张量核心加速）。

低精度注意力

越来越多的工作表明，KV-Cache 可以量化到低位，而精度损失可忽略不计。FlashInfer 实现了高性能的 fp8 解码核函数，与 fp16 核函数相比，可以将核函数加速高达2倍

图11：FP8 解码注意力性能，使用 Llama2-7B 设置：num_kv_heads=num_qo_heads=32，head_dim=128。序列长度从32到65536不等。

fp8 和 fp16 核函数的带宽利用率之间存在一些差距，但随着查询长度的增加，差距正在缩小。

Atom 在 FlashInfer 之上实现了带有 int4 量化的高性能解码注意力核函数。

页面大小对 FlashInfer PageAttention 的影响

FlashInfer 解码核函数在 GPU 共享内存中预取页面索引，从而最大限度地减少页面大小对核函数性能的影响。下面是 FlashInfer PageAttention 在 A100 上使用不同页面大小的性能比较

图12：不同 page_size 下的批量解码性能。batch_size 设置为1，使用 Llama2-7B 设置：num_kv_heads=num_qo_heads=32，head_dim=128。序列长度从32到65536不等。我们还添加了一条参考线，表示不使用页表的 FlashInfer 单请求解码注意力的性能。

4种不同页面大小的内存带宽利用率几乎相同，并且接近单请求解码注意力曲线，这表明页面大小对 FlashInfer PageAttention 的核函数性能影响很小，且页表本身的开销很小。

一些最近的工作，例如 sglang，探索了需要 page_size=1 的新型 KV-Cache 管理算法，性能可以从 FlashInfer 的优化中受益。

总结与未来工作

在序列维度上拆分 KV-Cache 以增加并行度的想法也在 Flash-Decoding 中进行了探索，FlashInfer 同时实现了这一想法，请参阅我们的2023年9月1日的 github 检查点和我们于2023年9月5日在 TVM Unity 开放开发会议上的公开演讲。

目前 FlashInfer 仅支持 NVIDIA GPU，FlashInfer 的 AMD 和 Apple GPU 版本已在 MLC-LLM 项目中得到初步支持，这得益于 Apache TVM 编译器。我们的下一个版本将包括 Atom 中提出的4位融合解量化+注意力操作符以及 Punica 中使用的 LoRA 操作符。从长远来看，我们对 post-Hopper NVIDIA GPU 和 AMD/Apple GPU 上的性能优化以及新兴 LLM 架构中的新操作符感兴趣。请查看我们的路线图了解开发计划，并留下您希望在 FlashInfer 中看到的功能建议。

致谢

FlashInfer 的灵感来源于 FlashAttention 2、vLLM、cutlass 和 Stream-K 项目。

这篇博客文章由 Zihao Ye撰写。我们感谢整个 FlashInfer 团队对该项目的贡献

• Zihao Ye (UW)：FlashInfer 的设计与实现
• Lequn Chen (UW)：页表数据结构设计、API 设计、CI/CD 和 Punica 集成
• Ruihang Lai (CMU)：KV-Cache 设计、API 设计以及与 MLC-LLM 的集成
• Yilong Zhao (UW & SJTU)：int4 注意力操作符
• Size Zheng (UW & PKU)：CUDA 优化和推测性解码
• Junru Shao 和 Yaxing Cai (OctoAI)：MLC-LLM 集成
• Bohan Hou 和 Hongyi Jin (CMU)：使用 Apache TVM 将 FlashInfer 移植到 AMD 和 Mac GPU
• Liangsheng Yin (SJTU & LMSys)：PyTorch 绑定和 sglang 集成。
• Yifei Zuo (UW & USTC)：PyTorch 绑定
• Tianqi Chen (CMU & OctoAI)：softmax/注意力合并的递归形式和建议
• Luis Ceze (UW & OctoAI)：性能分解分析和建议

我们还要感谢 Masahiro Masuda (OctoAI)、Yixin Dong (UW & SJTU)、Roy Lu (UW)、Chien-Yu Lin (UW)、Ying Sheng (Stanford & LMSys) 和 Lianmin Zheng (Berkeley & LMSys) 的宝贵反馈和讨论。

引用

@misc{flashinfer,
    title = {Accelerating Self-Attentions for LLM Serving with FlashInfer},
    url = {https://flashinfer.cn/2024/02/02/introduce-flashinfer.html},
    author = {Ye, Zihao and Chen, Lequn and Lai, Ruihang and Zhao, Yilong and Zheng, Size and Shao, Junru and Hou, Bohan and Jin, Hongyi and Zuo, Yifei and Yin, Liangsheng and Chen, Tianqi and Ceze, Luis},
    month = {February},
    year = {2024}
}

脚注

1. Dissecting Batching Effects in GPT Inference by Lequn Chen ↩