CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection¶

会议: NeurIPS 2026
arXiv: 2605.16839
代码: https://github.com/jiwonsong-dev/CompactAttention
领域: LLM效率
关键词: 分块预填充, 稀疏注意力, KV Cache 管理, 分页注意力, 分组查询注意力

一句话总结¶

CompactAttention 针对大语言模型分块预填充（Chunked Prefill）中查询远短于上下文导致的稀疏内核低效与显式 KV 拷贝开销，提出将 2D 块稀疏掩码通过查询块并集与组内并集降阶为 GQA 友好的 KV 块表，并结合 KV-head-major 内存布局实现就地零拷贝分页注意力加速。

研究背景与动机¶

随着大语言模型在长程推理、文档理解和代码分析中的广泛应用，上下文窗口已经扩展到数十万甚至数百万 token。在在线服务系统中，单次完成完整长上下文预填充（One-shot Prefill）不仅计算量呈二次方爆炸，还会长期占用计算资源，导致批处理中的解码请求严重停顿，无法满足 token 间生成延迟（TBT）的服务等级目标。为此，vLLM、SGLang 等现代推理系统普遍采用分块预填充（Chunked Prefill），将长输入拆分为固定大小的计算块逐步处理，每次计算仅让当前查询块与历史累积的 KV Cache 发生交互。

然而，现有的长上下文注意力加速机制直接移植到分块预填充时遭遇了根本瓶颈。传统块稀疏注意力（如 MInference、FlexPrefill、XAttention）专为单次全量预填充设计，在分块预填充场景下面临严重的短查询非对称性（$Q \ll KV$）：由于每次迭代中活跃的查询块极少，块稀疏 CUDA 内核不仅缺乏足够的计算并行度，而且解释稀疏掩码和非规则内存访问的固定开销无法被分摊，导致实际加速远低于理论稀疏度收益；同时，跨 chunk 累积的在线模式搜索也会造成巨大的计算浪费。另一方面，直接针对分块预填充设计的 QUOKA 尝试通过查询子采样选出关键 KV 并执行密集计算，但子采样会遗漏非采样 query 独有的局部关键 KV，在多键检索（Multi-key NIAH）等分布密集型任务上导致准确率显著下降，且其 token 粒度筛选需要执行显式的 KV 聚合搬移，内存带宽开销随上下文线性激增。

面对这一两难困境，核心矛盾在于稀疏执行内核在短查询下的低效性与 token 级筛选显式内存搬移开销之间的冲突。本文切入点是将“稀疏 KV 块的选择逻辑”与“底层注意力的执行机制”解耦：不再依赖低效的稀疏执行内核，而是利用成熟轻量的稀疏模式预测器筛选出关键 KV 块，并将其映射为紧凑的分页查找表。核心 idea：通过查询块并集（Q-block Union）与组内并集（Intra-group Union）将 2D 块稀疏掩码降阶为覆盖保全的 GQA 分组 KV 块表，配合 KV-head-major 缓存布局驱动原生分页注意力内核实现就地零拷贝计算。

方法详解¶

整体框架¶

CompactAttention 的整体流水线分为选择阶段（KV Selection）与执行阶段（Zero-Copy Paged Execution）。在每个分块预填充迭代中，输入当前 chunk 的查询块与之前累积的 KV 缓存。选择阶段接收轻量级块稀疏模式预测器生成的 2D 稀疏掩码，依次通过两级布尔并集操作将其收敛为按 GQA 执行组划分的 1D KV 块表；执行阶段将该元数据转换为 CSR 格式的页索引指针，直接送入经过优化的高性能稠密分页注意力后端（如 FlashInfer），在不搬移任何 KV 张量物理内存的前提下原地完成注意力计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入：当前 Chunk 查询与累积 KV 缓存"] --> B["轻量 2D 块稀疏模式预测<br/>SeerAttention / FlashPrefill"]
    B --> C["块并集 KV 表构建<br/>Q 块并集压缩 + 组内并集映射"]
    C --> D["零拷贝分页注意力执行<br/>KV-head-major 布局 + FlashInfer 原地计算"]
    D --> E["输出：当前 Chunk 注意力输出并追加 KV"]

关键设计¶

1. 块并集 KV 表构建：将二维细粒度稀疏掩码转换为保全覆盖的分组页表

传统块稀疏注意力将二维布尔掩码 $M_{b,h,i,j} \in \{0, 1\}$（对应批次 $b$、头 $h$、查询块 $i$ 与 KV 块 $j$）直接作为稀疏算子的执行调度图，在 $Q \ll KV$ 时陷入硬件利用率极低的困境。CompactAttention 转变思路，将二维掩码视为候选 KV 块的检出信号。由于标准密集分页注意力内核对同一个执行组内的所有查询块采用统一的 KV 访问列表，CompactAttention 首先对查询块维度 $i$ 执行布尔或操作（Q-block Union）： $$ \bar{M}{b,h,j} = \bigvee $$ 这消除了查询块维度的非规则跳跃，得到每个头专属的 1D KV 块掩码。随后，针对分组查询注意力（GQA）中多个查询头共享同一组 KV 头的拓扑结构，CompactAttention 进一步在共享 KV 块表的查询头集合 } M_{b,h,i,j$\mathcal{H}(g)$ 内执行组内并集（Intra-group Union）： $$ G_{b,g,j} = \bigvee_{h \in \mathcal{H}(g)} \bar{M}_{b,h,j} $$ 最终得到每个执行组 $g$ 的紧凑页表 $T_{b,g} = \{j \mid G_{b,g,j} = 1\}$。该两级并集操作在数学上构成了满足分页注意力共享约束下的最小完全覆盖集（Minimal Coverage-Preserving Set）：只要原始 2D 掩码中任意查询块或查询头命中了块 $j$，该块就会被完整保留在页表中，从而彻底规避了 QUOKA 等子采样方法漏选特定 query 关键上下文导致的精度退化。虽然并集操作相比理论 2D 块稀疏降低了一定稀疏度，但由于底层的密集分页内核执行效率远高于非规则稀疏内核，这一工程置换换取了显著的端到端吞吐提升。

2. 零拷贝分页注意力执行：基于 KV-head-major 内存布局的原地分页访问

选出离散的 KV 块后，传统基于候选缩减的方法通常需要将选中的 KV 向量重新搬移（gather-and-pack）到一块连续的紧凑缓冲区内，这种显式内存拷贝开销随着上下文长度和批大小的增加急剧膨胀，完全抵消了计算量缩减带来的加速。CompactAttention 提出通过对底层 KV Cache 物理布局的协同改造来消除数据拷贝。常规框架采用序列主序（Sequence-major）布局 $[B, L, H_{kv}, D]$，使得同一个 token 的所有 KV 头在物理内存中紧密交织，强制所有注意力头共享相同的页分配表，无法支持各个 GQA 组独立筛选块。

CompactAttention 采用 KV 头主序（KV-head-major）布局 $[B, H_{kv}, L, D]$，将每个批次、每个 KV 头的连续块在内存中组织为独立的连续物理区域 $[block\_size, D]$。配合定制的融合 CUDA 内核，CompactAttention 仅需在 CPU/GPU 控制侧将并集掩码快速转化为 CSR 稀疏格式的指针数组（kv_indptr 与 kv_indices），将批次与 KV 头维度平铺为伪批次序列，以 num_kv_heads=1 直接调用 FlashInfer 的 Paged Attention 接口。底层的 Tensor Core 原生利用页表指针做带跨距的就地读取，彻底避免了任何 KV 张量的显式分配与数据拷贝。此外，当前正在计算的最新 chunk 始终在掩码中无条件全开，确保因果自注意力（Causal Masking）在物理真实绝对位置上严谨成立，避免压缩坐标系下的语义破损。

3. 子 KV 组划分（Sub-KV-Group Union）：抑制大 GQA 组下的稀疏度退化

在 Qwen 等采用极端非对称 GQA 架构的模型中（如查询头与 KV 头比例达到 8:1），若直接对整个 KV 组内的 8 个查询头做并集，过多注意力头模式的叠加会导致保留比例（Keep ratio）急剧上升，使模型退化为近乎全量稠密计算。为了兼顾计算并行度与稀疏收益，CompactAttention 提出了子 KV 组划分策略：将每个大 KV 组细切为包含固定数量（默认设为 4 个）查询头的子执行组，在子组内部独立执行组内并集并维护专属的分页块表。子组划分有效控制了并集引入的稀疏度稀释，在维持零拷贝硬件接口的前提下，使超长上下文下的注意力计算保持高度精简。

损失函数 / 训练策略¶

CompactAttention 属于无需重训练或微调的推理时计算加速范式（Training-free / Plug-and-play）。在模式预测器选择上，它解耦且兼容学习型（如 SeerAttention 的轻量门控网络 AttnGate）与免训练型（如 FlashPrefill 基于动态阈值 $\alpha$ 的最大值截断法）两种机制。为了补偿两级并集带来的稀疏度回收，CompactAttention 针对初始 2D 掩码预测采用更为激进的截断超参数（例如在 LLaMA-3.1-8B 上将 FlashPrefill 的判定阈值从 $\alpha=0.01$ 收紧至 $\alpha=0.06$），使得经过块并集操作后的实际物理计算稀疏度与原始 FlashPrefill 保持相当（约 70% 稀疏），在无损保留长文本下游精度的同时最大化硬件执行效率。

实验关键数据¶

主实验¶

评估在分块预填充配置下进行，评测基准涵盖长文本基准 RULER（检索与聚合）与 LongBench V2（深度长上下文理解），基座模型采用 LLaMA-3.1-8B-Instruct（128K 上下文）与 Qwen3-30B-A3B-Instruct-2507（MoE，256K 上下文）。测试硬件包括 RTX PRO 6000（Blackwell 架构，GDDR7，FlashAttention-2）与 H200 SXM（Hopper 架构，HBM3e，FlashAttention-3）。

下表展示了在分块预填充（Chunk Size 1024）下 RULER 长上下文任务在不同长度下的准确率表现（摘自原文 Table 1）：

方法	模型	32K 准确率 (%)	64K 准确率 (%)	128K 准确率 (%)	平均准确率 (%)
Dense（稠密基线）	LLaMA-3.1-8B	88.23	83.52	76.59	82.78
QUOKA	LLaMA-3.1-8B	83.52	79.99	70.44	77.98
XAttention	LLaMA-3.1-8B	89.16	83.27	74.50	82.31
SeerAttention	LLaMA-3.1-8B	90.23	83.82	73.34	82.46
CompactAttention-SA (Ours)	LLaMA-3.1-8B	88.92	84.12	74.28	82.44
FlashPrefill	LLaMA-3.1-8B	88.99	83.94	74.12	82.35
CompactAttention-FP (Ours)	LLaMA-3.1-8B	88.77	83.96	74.17	82.30
Dense（稠密基线）	Qwen3-30B-A3B	93.01	90.56	87.81	90.46
QUOKA	Qwen3-30B-A3B	81.96	79.57	79.03	80.19
XAttention	Qwen3-30B-A3B	92.29	88.92	88.13	89.78
FlashPrefill	Qwen3-30B-A3B	91.55	87.73	86.77	88.68
CompactAttention-FP (Ours)	Qwen3-30B-A3B	92.41	88.43	87.26	89.37

下表展示了 LLaMA-3.1-8B-Instruct 在 H200 SXM 平台（TP=2, Batch Size 8, Chunk Size 1024）上的分块预填充端到端及注意力延迟对比（摘自原文 Table 4）：

方法	8K 注意力 (ms)	32K 注意力 (ms)	128K 注意力 (ms)	128K 端到端 (ms)	128K 端到端加速比
Dense	438.6	4225.4	55619.7	70754.4	1.00×
QUOKA	582.3	4629.7	51292.4	66029.1	1.07×
XAttention	923.9	7729.7	86821.4	103870.2	0.68×
SeerAttention	1002.5	9478.1	85852.0	105207.6	0.67×
FlashPrefill	474.1	3494.6	24876.1	41161.7	1.72×
CompactAttention-SA (Ours)	780.2	4393.8	29468.1	45366.8	1.56×
CompactAttention-FP (Ours)	533.6	3160.1	19899.9	35806.2	1.98×

消融实验¶

在 Qwen3-30B-A3B（GQA 比例 8:1）上的子组大小（Subgroup Size）消融分析（摘自原文 Table 2，H200 SXM, TP=2, Batch Size 4, 128K 上下文, Chunk Size 1024, $\alpha=0.12$）：

配置	子组大小	块保留比 (Keep ratio)	RULER 128K 准确率 (%)	注意力延迟 (ms)	说明
激进子组	2	0.200	87.03	14790.7	稀疏度最高，延迟最低，精度轻微受损
默认子组 (Ours)	4	0.280	87.26	16514.6	精度与吞吐的最佳平衡折中
全 KV 组并集	8	0.372	87.60	19732.4	8个头并集导致过多块保留，计算变慢

在执行机制消融（原文 Figure 7(c) 对应分析）中，对比固定并集掩码下的三种执行策略： 1. 块稀疏内核执行：直接使用稀疏注意力算子，由于短查询缺少并行度，延迟最高。 2. 显式拷贝聚合执行（CompactAttention-FP Copy）：将选出的离散 KV 块搬移到连续缓存区再跑密集计算，内存带宽与拷贝开销抵消大量收益。 3. 原地零拷贝分页执行（CompactAttention-FP 原生）：仅构建轻量 CSR 元数据后直接驱动 FlashInfer，取得了显著最低的计算延迟。

关键发现¶

稀疏算法的执行脱节是核心瓶颈：XAttention 和 SeerAttention 等先进稀疏算法在单次全量预填充表现优异，但在分块预填充中因每个 chunk 重复搜寻的开销以及非对称块稀疏算子的低效，端到端延迟甚至大幅落后于 Dense 基线（在 128K 下比 Dense 慢近 50%）。
零拷贝分页执行碾压显式数据搬移：相比于 QUOKA 在 128K 下仅取得 1.07× 的微弱端到端加速，CompactAttention-FP 依靠原地零拷贝机制消除数据搬运，在 128K 下达到了 2.79× 的注意力加速与 1.98× 的端到端加速。
全查询块评估有效防止局部信息丢失：QUOKA 因查询采样在 RULER 的多键提取任务（Multi-key NIAH-3）上断崖式下跌 7.4%，而 CompactAttention 凭借并集保全性，在各尺度下与全精度 Dense 基线的平均差距控制在 0.5% 以内。

亮点与洞察¶

将稀疏掩码从“执行算子”重新定义为“分页路由元数据”：以往稀疏注意力一直执着于定制极致优化的稀疏 CUDA 内核，CompactAttention 洞察到工业界成熟的 PagedAttention 已经高度利用了 GPU 张量核心与内存流水线，将稀疏问题转化为页表路由元数据构建，巧妙将算法稀疏度嫁接到了成熟的工业级稠密推理引擎上。
覆盖保全的两级并集抵消了粗粒度带来的信息丢失：通过 Q-block union 覆盖所有查询可能关注的上下文，通过 intra-group union 契合硬件 GQA 约束，虽牺牲了少许理论稀疏度，却换取了近乎无损的下游检索精度与极高的算力利用率。
软硬件系统协同设计：针对性地采用 KV-head-major 内存布局取代序列主序，使得以 KV 块为单位的动态元数据索引无需产生跨头数据交织，是实现 zero-copy 的关键底层支撑。

局限与展望¶

依赖上游模式搜索器的召回质量：CompactAttention 本身是降阶与执行框架，其精度上限受制于前端块稀疏预测器（如 SeerAttention 或 FlashPrefill）的漏检率，一旦前端掩码遗漏了关键信息块，后端的并集操作无法恢复。
短上下文与小批次下元数据构建开销无法分摊：在短序列（如 $\le 8\text{K}$）或极小批次下，累积 KV 较小，注意力占整机耗时比例低，此时构建页表的 CPU/GPU 控制开销与模式搜索开销可能超过注意力本身的缩减收益，导致比稠密基线更慢。
服务系统侵入性改造门槛：该方案需要修改现有 serving 框架（如 vLLM、SGLang）底层的 KV Cache 分配策略，将其从 sequence-major 改造为 KV-head-major，并开放针对每个执行组下发差异化 block_table 的底层算子接口。

评分¶

新颖性: ⭐⭐⭐⭐☆ 突破了传统必须使用块稀疏算子执行稀疏注意力的固有思维，提出用掩码并集构建分页块表的降阶路线。
实验充分度: ⭐⭐⭐⭐⭐ 跨 Blackwell 与 Hopper 两种架构硬件，对比了全套块稀疏与预填充候选方案，覆盖完整延迟与消融。
写作质量: ⭐⭐⭐⭐⭐ 逻辑严密，图文对照精确，对短查询下硬件瓶颈与布局机制的剖析非常透彻。
价值: ⭐⭐⭐⭐⭐ 完美契合当今大模型推理服务中分块预填充成为标配的工业现实，具有极高的工程落地与部署价值。