💬 NLP

CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

CompactAttention 通过引入一种块并集键值选择机制，将二维块稀疏掩码转换为高效的、支持分组查询感知的每组键值块表，从而在保持接近稠密精度的同时实现无需显式压缩的原地内存访问，并将长上下文大语言模型中的分块预填充速度提升高达 2.72 倍。

原作者： Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim

发布于 2026-05-19

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你是一位图书管理员（即人工智能），试图基于一座庞大的图书馆（即上下文）来回答问题。在过去，如果你拥有一座巨大的图书馆，你就必须阅读每一本书才能找到正确的答案，这耗时极长。为了加快这一过程，研究人员发明了一种“分块预填充”系统：你不再一次性阅读整座图书馆，而是将其分成小批次（块）来阅读，并在阅读过程中将笔记添加到一个记事本（即 KV 缓存）中。

然而，一个新的问题出现了：当你收到新的一批问题时，如何快速在记事本中找到正确的页面，而无需每次都重新阅读所有内容？

本文介绍了CompactAttention，一种解决该问题的新方法。以下是其工作原理，使用简单的类比进行说明：

旧方法的问题

本文指出了人们曾尝试解决该问题的两种主要方式，以及它们失败的原因：

“稀疏内核”方法（低效的扫描仪）：
- 核心思路： 想象你拥有一张图书馆地图，上面用红点标记了仅重要的书籍。你试图跳过空白区域，只查看红点。
- 失败原因： 当你阅读一座巨大的图书馆（长上下文）但只提出一个微小的问题（小块）时，这种方法会变慢。这就像拥有一台擅长扫描整面文字墙的扫描仪，但当你只有一句话时，扫描仪的设置和校准时间过长。跳过空白区域的“开销”实际上使其比直接阅读所有内容还要慢。
“查询子采样”方法（懒惰的图书管理员）：
- 核心思路： 与其检查每一个问题，不如从你的批次中随机挑选几个问题，找出对这些问题重要的书籍，并假设这些书籍对所有人也同样重要。
- 失败原因： 这很冒险。如果你选错了那几个问题，可能会漏掉某本至关重要的书，而这本书只有某一个特定问题才需要。此外，一旦你选定了这些书，你就必须将它们从书架上物理搬运到一张特殊桌子上才能阅读。这种“搬运”（复制数据）需要耗费大量时间和能量。

解决方案：CompactAttention

CompactAttention 通过将寻找书籍与阅读书籍分离开来，改变了游戏规则。

步骤 1：“并集”策略（分组搜索）

CompactAttention 不使用复杂的“跳表”（稀疏内核）或基于少数问题进行猜测，而是采用了一种聪明的分组技巧：

想象你有一个侦探团队（查询头）正在处理一个案件。每位侦探都有他们自己认为重要的“嫌疑人”（KV 块）列表。
CompactAttention 没有让每位侦探单独工作，而是说：“让我们把整个团队的所有嫌疑人合并成一份主列表。”
这分两步完成：
1. Q-Block 并集： 合并当前批次中所有问题的列表。
2. 组内并集： 合并协同工作的侦探们的列表。
结果： 你得到了一份单一的、最小化的“嫌疑人主列表”，涵盖了所有人的需求。没有任何重要的书会被遗漏，因为只要任何一位侦探需要它，它就会在列表中。

步骤 2：“零拷贝”执行（原地读取）

这是神奇的部分。

旧方式： 一旦你有了主列表，你就必须将所有这些书从书架物理移动到一张特殊桌子上，以便快速阅读。这种“移动”需要时间。
CompactAttention 方式： 你根本不需要移动这些书。你只需交给图书管理员一张地图（元数据），上面写着：“去 A 架第 3 排第 5 本书；然后去 B 架第 1 排第 2 本书。”
图书管理员（计算机内核）直接前往书架上的这些位置并读取它们。这被称为**“零拷贝分页注意力”**。它节省了所有在移动数据上花费的时间和能量。

为何意义重大

本文在具有 128,000 个词（非常长的文档）上下文的巨大 AI 模型（LLaMA-3.1-8B）上测试了这种方法。

准确性： 它与阅读整座图书馆（稠密注意力）一样聪明。它没有遗漏任何关键细节。
速度： 它比标准方法快高达2.72 倍。

核心启示

将 CompactAttention 想象成一位聪明的图书管理员，他不再试图重新整理图书馆，而是直接使用一张完美的、合并后的索引卡。

通过认识到“搜索”（寻找重要块）和“执行”（读取它们）应当分离，并利用“分组”技巧确保不遗漏任何内容，他们成功地在没有损失任何智能的情况下，显著加快了长文档 AI 处理的速度。他们证明了瓶颈不仅仅在于选择哪些书，更在于你如何去选取它们。

技术摘要：CompactAttention

问题陈述
随着大语言模型（LLM）扩展以支持数十万至数百万 token 的上下文窗口，由于二次方计算成本以及在在线服务系统中导致解码请求停滞的风险，通过单次“一次性”预填充（prefill）处理长序列已变得不切实际。因此，分块预填充（chunked prefill）——即按固定大小的块顺序处理输入，同时关注累积的键值（KV）缓存——已成为 vLLM 和 SGLang 等框架中的标准服务策略。

然而，现有的加速方法在该场景下面临显著局限：

块稀疏注意力低效：传统的块稀疏方法（如 XAttention、SeerAttention）依赖于稀疏核，这些核需要较长的查询长度来分摊不规则内存访问和掩码解析的开销。在分块预填充中，查询长度受块大小限制（ $Q \ll KV$ ），导致这些核失去并行性和效率，往往无法超越稠密注意力。
模式搜索开销：稀疏方法需要查找依赖于输入的稀疏模式。在分块预填充中，这种搜索必须针对每个块在日益增长的累积 KV 缓存上重复进行，使得累积搜索开销成为关键瓶颈。
直接 KV 选择的局限性：像 QUOKA 这样的方法通过基于子采样查询集选择 KV 子集并执行稠密注意力来避免稀疏核。然而，该方法存在两个问题：
- 覆盖缺口：对未采样查询重要的特定 KV 条目可能被遗漏，导致在需要分布式信息访问的任务上精度下降。
- 复制开销：Token 级选择需要在执行注意力之前，显式地收集并将 KV 条目复制到一个紧凑缓冲区中，引入的开销随上下文长度和批次大小而增加。

方法：CompactAttention
作者提出了 CompactAttention，这是一种将块级 KV 选择与注意力执行解耦的机制。CompactAttention 不使用二维块稀疏掩码作为稀疏核的直接执行计划，而是将其视为 KV 选择信号，以构建 GQA 感知的每组 KV 块表。

该方法分两个阶段运行：

选择阶段（块并集 KV 表构建）：
- 系统接受由任何轻量级模式搜索方法（如 SeerAttention 或 FlashPrefill）生成的二维每头块稀疏掩码。
- 它对当前块内的查询块应用 Q-块并集。这是必要的，因为稠密分页注意力为一起执行的一组查询块消耗单个 KV 块列表，而不是为每个 $(Q\text{-block}, KV\text{-block})$ 瓦片进行单独决策。
- 它对共享 KV 组的查询头（默认为分组查询注意力组）应用 组内并集。
- 结果是一个最小的每组 KV 块表（ $T_{b,g}$ ），保留了原始输入掩码选择的所有 KV 块。该构建确保了覆盖保留：组内任何查询块或头选择的任何 KV 块都保持被选中。
执行阶段（零拷贝分页注意力）：
- 与 Token 级选择方法不同，CompactAttention 在 块粒度 上运行。
- 它利用 KV-头主内存布局 $[B, H_{kv}, L, D]$ ，而不是标准的序列主布局。这使得每个 $(batch, KV\text{-head}, block)$ 三元组可以被视为连续的内存区域。
- 构建的块表作为 元数据（具体为 kv_indptr 和 kv_indices）传递给标准的 分页稠密注意力核。
- 该核从原始 KV 缓存中就地访问所选 KV 块，无需显式复制或压缩到新缓冲区。

主要贡献

解耦选择与执行：本文提出了一种设计，将 选择哪些 KV 块的逻辑与 如何执行 它们分离，允许重用轻量级稀疏模式搜索方法，同时避免稀疏核的低效性。
块并集构建：创新性地使用 Q-块和组内并集，将二维稀疏掩码转换为可执行的、GQA 感知的块表，这些表既最小化又保留了覆盖范围。
零拷贝分页执行：通过利用 KV-头主布局和张量元数据，该方法消除了像 QUOKA 这样的 Token 级选择方法固有的显式 KV 复制开销。

实验结果
在 LLaMA-3.1-8B-Instruct 和 Qwen3-30B-A3B-Instruct-2507 上进行分块预填充评估：

加速比：在 H200 GPU 上，针对 128K 上下文长度，CompactAttention 相比稠密注意力基线实现了高达 2.72 倍的注意力加速 和 1.96 倍的端到端加速。这显著优于块稀疏基线（在 $Q \ll KV$ 场景下通常表现不如稠密注意力）和 QUOKA（受限于复制开销）。
精度：在 RULER 和 LongBench V2 基准测试中，CompactAttention 保持了接近稠密注意力的精度。它优于 QUOKA，特别是在需要分布式信息访问的任务上（例如多键 NIAH），在这些任务中 QUOKA 的查询子采样会导致精度下降。
效率：消融研究证实，即使考虑元数据构建开销，零拷贝分页执行策略也比稀疏核执行和显式 KV 压缩更高效。

意义
本文指出，分块预填充中稀疏注意力的主要瓶颈不仅仅是 选择哪些 KV 块，而是 如何执行 所选块。通过块并集构建和零拷贝执行，使稀疏选择与高效的分页稠密注意力核兼容，CompactAttention 为实现更快的长上下文 LLM 服务提供了一条实用路径。它解决了阻碍现有稀疏注意力方法在分块预填充系统中采用的 $Q \ll KV$ 场景下的特定低效问题。