论文标题：Progressive Sparse Attention: Algorithm and System Co-design for Efficient Attention in LLM Serving

论文链接：https://arxiv.org/pdf/2503.00392

最近 DeepSeek 公开的 NSA 论文和 Kimi 公开的 MoBA 论文使得动态稀疏注意力机制广受关注，我和我的实习生 CUHK 周启辉等人最近半年也致力于动态稀疏注意力的研究，并完成了渐进稀疏注意力PSA (Progressive Sparse Attention) 的工作，分享给大家，希望和大家共同探讨该技术。

我们提出的 PSA 机制解决的问题是，现有动态稀疏注意力算法（包括 NSA 和 MoBA）中普遍采用的 Top-k 选择策略所存在的痛点问题，解决了该问题就清除了动态稀疏注意力技术走向商用的重大障碍。所以，理论上 PSA 是可以与 NSA 和 MoBA 完美结合实现更高的效率。

TL; DR

处理长上下文已成为现代大规模语言模型（LLM）的关键能力。然而，由于键值（KV）缓存占用大量显存，长上下文 LLM 推理的计算成本极高。现有研究利用动态稀疏注意力算法（DSA）来减少 KV 缓存的开销，但这些算法依赖于 Top-k 的KV缓存选择策略，其在推理准确性和性能之间存在权衡窘境：较大的 k 值提高了准确性但降低了推理性能，而较小的 k 值提升了推理性能但损害了准确性。

为了打破这一窘境，我们提出渐进稀疏注意力（PSA）机制，它结合算法创新与系统协同设计，同时实现了高的推理准确性和更优的推理性能。PSA 算法根据实际的注意力权重分布，自适应地调整不同 Token 和不同层的 KV 缓存使用预算，而非使用固定的 k 值预算，从而在减少 KV 缓存使用的同时保证较高的准确性。此外，为了提高系统执行性能，我们引入流水线迭代执行，一种面向PSA算法的高效执行方案，降低CPU和GPU间交错执行和同步的开销。同时，我们还设计了一种统一 GPU 内存管理策略，针对不同模型层的非均匀内存需求来优化 PSA 的内存利用率。大量实验表明，与最先进的 DSA 方法和不使用稀疏注意力的系统相比，PSA 将注意力计算的 KV 缓存使用减少最多 2.4 倍和 8.8 倍，同时将端到端推理吞吐量提高最多 1.4 倍和 2.0 倍。

1 稀疏注意力机制

静态稀疏注意力：由于大模型的自回归生成特点，每生成一个 Token 都需要从 GPU HBM显存中加载全部的 KV 缓存到片上 SRAM 上。对于长上下文推理而言，庞大的 KV 缓存会导致显著的时间和空间开销。近期工作发现，大模型的注意力计算实际上呈现高度稀疏性，即只有少部分 Token 的KV 缓存对最终的注意力权重贡献大。基于这一观察，已有若干 KV 缓存驱逐算法如H2O、SteamingLLM等提出在生成过程中静态丢弃不重要 Token 的 KV 缓存，以维持较小的缓存大小。然而，这些被丢弃的 Token 可能对未来的生成依然重要，从而可能造成关键信息的丢失。

动态稀疏注意力：为了防止信息丢失，动态稀疏注意力算法（DSA）如InfLLM、Quest、ArkVale等应运而生。图 1 展示了现有 DSA 的工作流程。与直接丢弃不重要 Token的KV 缓存不同，这些算法会保留所有 KV 缓存，并为每个查询 Token 动态选取一小部分关键 KV 缓存来进行注意力计算。 为了加速选择过程，受 PagedAttention中基于块级内存分配的启发，DSA通常 将 KV 缓存划分为多个块，并在块级别上进行选择。对于每个 KV 块，DSA 构造元数据向量以代表块内的 Token。不同的 DSA 算法提出不同的元数据构造方法，从简单的 Token 键均值计算到寻找 Token 键的包围盒（bounding cuboid）均有涉及。无论采用何种方法，由于元数据体积远小于 KV 块，DSA 可以通过计算元数据向量与查询 Token (Q) 的点积，来估计各个 KV 块与Q 之间的注意力分数，并选择分数最高的 Top-k个KV 块进行近似注意力计算。

图1：动态稀疏注意力工作流

2 Top-k选择的困境

现有 DSA 算法依赖 Top-k 选择策略，即在所有层和所有查询 Token 上均采用相同的 KV 缓存预算（即 k 值）进行注意力计算。为了确保高推理准确性，被选出的 Top-k 个KV 块应当与查询 Token的Q 具有足够高的累积注意力权重（例如 95%）。然而，我们观察到，由于不同 Token 之间和不同层之间的稀疏性都存在非常大的差异，固定的 Top-k 选择策略往往难以一致地达到这一要求：

（1）Token 间的注意力稀疏性存在很大差异：我们使用 LongBench 中两个热门数据集 QMSum 和 GovReport 评估了注意力稀疏性的差异，如图 2 所示。可以观察到，为了达到 95% 的累计注意力权重，不同查询 Token 所需要的 KV 块数量存在显著差异。例如，在 GovReport 的第 32 层（L32）中，20% 的查询 Token 需要少于 50 个 KV 块，60% 的 Token 需要 50 至 100 个 KV 块，而剩余 20% 的 Token 则需要超过 100 个 KV 块。

（2）不同层之间的注意力稀疏性也存在很大差异：我们同样观察到，不同层为了达到相同的累计注意力权重所需的 KV 块数量存在显著差异。例如，在 QMSum 的第 9 层（L9），80% 的查询 Token 需要少于 50 个 KV 块，而在第 14 层（L14）这一比例下降至 40%。

因此，很难为所有 Token 设定一个统一的 KV 块预算。较小的 k 值可能导致大多数 Token 的累计注意力权重不足，从而影响准确性；而较大的 k 值虽然可以满足准确性要求，但会导致选取过多 KV 块，影响系统效率。

图2：不同查询Token在不同的层实现95%累计注意力权重所需要的KV块数量的累积分布率CDF

3 渐进稀疏注意力PSA方法

3.1 System Overview

图3：基于PSA机制的推理系统架构

图 3 展示了基于 PSA 的推理系统 的整体架构，包含三个关键组件：批量调度器、模型执行器和 KV 缓存管理器。

3.2 PSA Algorithm

PSA 采用基于阈值的 KV 块选择算法，而不是基于 top-k 的选择，从而根据每个查询令牌的每一层的注意力权重分布动态分配 KV 块预算。具体来说，PSA 在 KV 块级别逐步执行注意力计算，并持续监控已计算 KV 块的累积注意力权重。一旦累积注意力权重达到给定的阈值（例如 95%），查询令牌的注意力计算立即终止，从而在实现了足够的准确率的同时，最小化KV 缓存的加载和计算。

算法1概述了 PSA 算法。给定查询向量 q 和解码请求的 KV 块索引 B，PSA 首先获取这些 KV 块的元数据 B_{meta}，并用它来计算每个 KV 块对查询的关键性分数 CS（第5行），类似于现有的 DSA。然后，PSA 根据关键性分数对这些 KV 块进行排序（第6行），并从最重要到最不重要的 KV 块开始执行注意力计算（第8-10行）。

PSA 在microbatch级别逐步执行注意力计算，其中一个microbatch次包含 m 个 KV 块—— m 是单次注意力计算所需的最小 KV 块数量，用于减少注意力内核调用的频率。在示例中，m 被设置为 1。每次注意力计算后，返回部分输出 O 和该块的注意力分数总和 AS（第10行）。PSA 然后将 O 合并到累积输出 O__{acc} 中（第11行），并更新累积的注意力分数总和 AS__{acc}（第12行）。

PSA 接着计算累积注意力分数 AS__{acc} 与总注意力分数 AS_{total} 的比例。然而，在对所有 KV 块执行注意力计算之前，AS_{total} 是未知的。为了解决这一问题，PSA 记录迄今为止遇到的最小块级注意力分数 AS_{min}（第13行），并将其用作尚未处理的剩余块的上界注意力分数，以估算 AS_{total}（第15行），计算公式为 AS_{acc} + AS_{min} * N_{left}。最终，可以使用 AS_{acc} 和 AS_{total} 来估算累积注意力权重 P_{acc}（第16行）。当 P__{acc} 超过阈值时，返回 O_{acc} 并完成注意力计算（第17-18行）。

算法1：PSA算法伪代码

图 4 展示了执行 PSA 算法的一个示例。给定一个按重要性分数排序的 KV 块列表和查询 Q，PSA 首先计算 Q 与第一个mirobatch的 KV 块（98、84、68、55）之间的注意力权重，得到累积注意力权重 P__{acc} = 0.61。由于 0.61 小于阈值ϵ （=0.98），PSA 继续计算 Q 与第二个microbatch的 KV 块（31、22、15、10）之间的注意力权重，得到 P_{acc} = 0.91。由于 0.91 仍然低于 ϵ，PSA 继续处理第三个microbatch KV 块（8.2、6.1、4.9、4.3），最终达到 P_{acc} = 0.98。当 P__{acc} 等于 ϵ 时，注意力计算完成。

图4：PSA算法执行示意图

3.3 Pipeline Iteration Execution

然而，将 PSA 算法实现在实际的 LLM 服务系统中面临着执行性能低的挑战。PSA 算法将注意力过程分解为多个迭代。在每个迭代中，PSA算法首先需要从主机内存中加載任何HBM中miss的KV块到 GPU HBM 中。然后，在准备好的 KV 块上执行标准注意力计算，再更新累积注意力权重，同时检查是否达到了预定义的阈值。这种动态的多迭代计算模式可能导致两个性能问题：（1）由于 CPU 数据准备和 GPU 计算的交错，导致 GPU 利用率低；（2）由于频繁从GPU传输注意力权重到CPU上来进行阈值判断，导致 CPU 和 GPU 之间的同步开销显著。

为了解决这些问题，我们提出了面向PSA 的流水迭代执行方案。具体来说，为了提高 GPU 利用率，我们将 KV 数据准备和注意力计算分配给不同的线程，确保数据加载和计算可以并发进行，而不会相互阻塞。为了进一步最大化并行性，我们为这些操作使用单独的 CUDA 流，允许数据传输和内核执行以流水线的方式重叠。这种设计最小化了 GPU 的空闲时间，并通过确保 GPU 持续用于计算，而数据准备在并行进行，从而提高了整体执行效率。

为了减少同步开销，我们设计了一个用与验证的 GPU 内核，该内核直接在 GPU 上更新和检查累积注意力权重，从而消除了将这些权重传输到CPU的需要。一旦累积注意力权重达到预定义的阈值，验证内核将使用零拷贝技术在固定主机内存中写入一个信号变量，通知 CPU 终止注意力过程。

3.4 Unified Memory Management

现有的 LLM 推理系统（例如 vllm）中的 GPU 内存管理通常是按层进行的。具体来说，系统首先通过一次预估执行确定可用的 GPU 内存容量，然后为每一层分配一个具有相等容量的 GPU 张量，并且各层之间相互独立使用。这种设计选择是基于现代 LLM 的 Transformer 层在内存使用模式上是相同的这一事实。也就是说，当执行一个批次的请求时，每个请求在所有层中访问的 KV 块数量是相同的。

然而，将这种按层分离的内存管理方法直接应用于 PSA 会导致 GPU 内存利用率低下。原因是 PSA 通过渐进注意力机制根据注意力权重分布动态调整涉及注意力计算的 KV 块数量。如图 2 所示，不同层的注意力权重偏斜度差异显著。那些偏斜度较低的层在注意力计算中访问的 KV 块数量远多于偏斜度较高的层，从而导致缓存命中率较低。

为了解决这一问题，我们提出了一种统一的 GPU 内存管理策略。具体来说，PSA 首先通过一次预估执行确定可用的 GPU 内存容量，然后分配一个具有该容量的单一 GPU 张量给所有层共享使用。该 GPU 张量随后被划分为等大小的槽，用于存储 KV 块，这些槽由 KV 缓存管理器进行管理。在运行时，所有涉及 KV 块的操作（包括分配、释放和加载）都发送到 KV 缓存管理器进行处理。

4 实验结果

4.1 KV Cache使用率

我们首先比较 vllm-sparse (top-k选择)、InfiniGen 和 PSA 在减少注意力计算中使用的 KV 缓存方面的有效性。我们确保平均请求准确性达到 98% 的目标，这是实际生产环境中常见的精度要求。

图5：不同数据集下vllm-sparse、InfiniGen和PSA平均KV Cache使用率

如图 5 所示，PSA 在减少 KV 缓存使用方面显著优于 vllm-sparse 和 InfiniGen。具体来说，与 vllm-sparse 相比，PSA 在 LWM-Text-7B 和 Llama-3.1-8B 上分别将 KV 缓存比例减少了 2.1 倍和 2.4 倍。这是因为 vllm-sparse 为所有请求分配统一的 KV 缓存预算，导致对短序列请求和注意力权重分布偏斜的请求过度选择 KV 块。相比之下，PSA 根据请求的注意力权重分布动态调整 KV 缓存预算，从而最小化 KV 缓存的使用。与 InfiniGen 相比，PSA 在 LWM-Text-7B 上将 KV 缓存比例减少了 1.8 倍。尽管 InfiniGen 使用细粒度的 token 级别选择，但它使用统一的最小注意力分数阈值来选择关键 Token，未能考虑不同请求中注意力分数范围的变化。此外，InfiniGen 采用权重矩阵压缩和层间预取来加速 KV 缓存的选择和加载过程，但这也无意中导致了注意力分数估计的精度下降。

4.2 端到端推理性能

我们测量了 vllm、vllm-sparse 和 PSA 在相同服务等级目标（SLO）要求下的请求吞吐量，如图 6 所示。根据之前的研究，我们将 P99 TBT的严格和宽松 SLO 分别定义为Decode迭代执行时间的 5 倍和 25 倍。

图6：不同SLO要求下，vllm, vllm-sparse和PSA的端到端QPS性能对比

与 vllm 相比，PSA 在严格 SLO 要求下分别将 LWM-Text-7B 和 Llama-3.1-8B 的吞吐量提高了 1.5 倍和 1.3 倍。在宽松 SLO 要求下，PSA 进一步将吞吐量分别提高了 2.0 倍和 1.5 倍。这是因为 vllm 的吞吐量受到其小推理批次大小的限制，而推理批次大小受到 GPU 内存的限制。相比之下，PSA 有效减少了参与注意力计算的 KV 缓存量，从而实现了更大的推理批次大小和更高的吞吐量。与 vllm-sparse 相比，PSA 在严格 SLO 要求下分别将 LWM-Text-7B 和 Llama-3.1-8B 的吞吐量提高了 1.3 倍和 1.2 倍。在宽松 SLO 要求下，PSA 分别将吞吐量提高了 1.4 倍和 1.3 倍。这是因为 PSA 通过渐进稀疏注意力机制比 vllm-sparse 减少了更多的 KV 缓存使用。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

第一阶段（10天）：初阶应用

该阶段让大家对大模型 AI有一个最前沿的认识，对大模型 AI 的理解超过 95% 的人，可以在相关讨论时发表高级、不跟风、又接地气的见解，别人只会和 AI 聊天，而你能调教 AI，并能用代码将大模型和业务衔接。

• 大模型 AI 能干什么？
• 大模型是怎样获得「智能」的？
• 用好 AI 的核心心法
• 大模型应用业务架构
• 大模型应用技术架构
• 代码示例：向 GPT-3.5 灌入新知识
• 提示工程的意义和核心思想
• Prompt 典型构成
• 指令调优方法论
• 思维链和思维树
• Prompt 攻击和防范
• …

第二阶段（30天）：高阶应用

该阶段我们正式进入大模型 AI 进阶实战学习，学会构造私有知识库，扩展 AI 的能力。快速开发一个完整的基于 agent 对话机器人。掌握功能最强的大模型开发框架，抓住最新的技术进展，适合 Python 和 JavaScript 程序员。

• 为什么要做 RAG
• 搭建一个简单的 ChatPDF
• 检索的基础概念
• 什么是向量表示（Embeddings）
• 向量数据库与向量检索
• 基于向量检索的 RAG
• 搭建 RAG 系统的扩展知识
• 混合检索与 RAG-Fusion 简介
• 向量模型本地部署
• …

第三阶段（30天）：模型训练

恭喜你，如果学到这里，你基本可以找到一份大模型 AI相关的工作，自己也能训练 GPT 了！通过微调，训练自己的垂直大模型，能独立训练开源多模态大模型，掌握更多技术方案。

到此为止，大概2个月的时间。你已经成为了一名“AI小子”。那么你还想往下探索吗？

• 为什么要做 RAG
• 什么是模型
• 什么是模型训练
• 求解器 & 损失函数简介
• 小实验2：手写一个简单的神经网络并训练它
• 什么是训练/预训练/微调/轻量化微调
• Transformer结构简介
• 轻量化微调
• 实验数据集的构建
• …

第四阶段（20天）：商业闭环

对全球大模型从性能、吞吐量、成本等方面有一定的认知，可以在云端和本地等多种环境下部署大模型，找到适合自己的项目/创业方向，做一名被 AI 武装的产品经理。

• 硬件选型
• 带你了解全球大模型
• 使用国产大模型服务
• 搭建 OpenAI 代理
• 热身：基于阿里云 PAI 部署 Stable Diffusion
• 在本地计算机运行大模型
• 大模型的私有化部署
• 基于 vLLM 部署大模型
• 案例：如何优雅地在阿里云私有部署开源大模型
• 部署一套开源 LLM 项目
• 内容安全
• 互联网信息服务算法备案
• …

学习是一个过程，只要学习就会有挑战。天道酬勤，你越努力，就会成为越优秀的自己。

如果你能在15天内完成所有的任务，那你堪称天才。然而，如果你能完成 60-70% 的内容，你就已经开始具备成为一名大模型 AI 的正确特征了。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】