CompactAttention：通过块联合 KV 选择加速分块预填充 CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection
用块联合 KV 选择替代稀疏注意力核，在 128K 上下文实现 2.72× 注意力加速。
前置知识
分块预填充 (Chunked Prefill)
vLLM、SGLang 等推理框架对超长输入的切分处理策略：把长 prompt 分成固定大小（如 512/1024 tokens）的 chunk 依次处理，每个 chunk 在累计 KV cache 时同时完成自身 Q 的注意力，避免一次性预填充阻塞 decode。
本文所有问题与方案都围绕这一场景：$Q \ll KV$ 的不平衡形态使现有稀疏注意力设计失效，方案也专门为此形态做执行优化。
块稀疏注意力 (Block-Sparse Attention)
把 Q/K 划分成 token block（如 64/128），先用轻量打分预测哪些 (Q-block, KV-block) tile 重要，再用稀疏内核只计算被选中的 tile，代表方法有 XAttention、SeerAttention、FlashPrefill，常见于一次性预填充加速。
作者把现有方法当作 2D 稀疏掩码产生器复用，但指出 block-sparse kernel 在分块预填充 $Q \ll KV$ 形态下并行度不足，提速远低于理论稀疏度。
分组查询注意力 (GQA) 与 Paged Attention
GQA 把多个 query head 共享同一组 KV head，降低 KV 存储与带宽；Paged Attention 借鉴操作系统虚拟内存的分页思想，把 KV cache 切成等长 page 并维护索引表，由 vLLM 推广为业界标准。
CompactAttention 的核心创新是构造 per-group KV page table，并用 KV-head-major 布局让每组拥有独立分块表，正是 GQA + Paged Attention 的组合产物。
研究动机
现有长上下文注意力加速在分块预填充下有两个根本瓶颈。其一是内核低效：在 $Q \ll KV$ 形态下（如 Q=1024、KV=128K），block-sparse kernel 的 query block 数量过少，无法摊销稀疏掩码解析与不规则访存的固定开销，Figure 1(b) 显示 KV=64K、稀疏度 90% 时一次性预填充接近理想 10×，而分块预填充下 XAttention、SeerAttention、FlashPrefill 都明显低于一次性预填充结果。其二是 pattern search 累积开销：每个 chunk 都要对累计 KV cache 重新打分，Figure 1(c) 表明 XAttention 在 64 个 chunk 上累计耗时已不可忽略。另一种思路 QUOKA 用 query 子采样 + token 级 KV 选择绕开稀疏 kernel，但 Figure 2 显示被忽略的非采样 query 可能漏掉 query-specific KV，且 token 级选择必须显式 gather 一次 KV，带来随上下文与 batch 增长的复制开销。
本文的目标是本文目标是设计一种分块预填充友好的注意力机制，在覆盖所有 query block、避免稀疏内核低效的同时，做到不引入显式 KV 复制开销。具体而言，希望保留 SeerAttention、FlashPrefill 等轻量级 block-level selector 的选择质量，但让 2D 稀疏掩码不再直接喂给稀疏 kernel，而是被转换成 GQA-aware 的每组 KV page table，由标准 paged dense-attention 后端在零拷贝下完成计算。目标是 128K 上下文、chunked prefill 配置下实现远超 XAttention、SeerAttention、FlashPrefill、QUOKA 的注意力与端到端加速，同时在 RULER 与 LongBench V2 两个长上下文基准上保持接近 dense 的精度。
与已有工作不同的是，本文的独特切入角度是「选择与执行解耦」——把稀疏注意力拆成两个独立子问题：哪些 KV block 该选（用现有轻量 selector 解决）以及如何高效计算这些被选中的 block（用 dense paged attention 解决）。这一视角与此前工作截然不同：XAttention、SeerAttention、FlashPrefill 沿用「2D 掩码 = 执行计划」范式，把选择信号直接当作 kernel 调度表；QUOKA 则把选择信号当成「要拷贝的 KV 索引」，都需要 token 级 gather。CompactAttention 首次把 2D block-sparse mask 仅当作选择信号，通过 Q-block union 与 intra-group union 把它 lower 成 GQA-aware per-group KV page table，同时享受 block 级 selector 的轻量与 dense paged kernel 的高效，并通过 KV-head-major 布局让每组拥有独立 page table，避免任何形式的 KV 复制。
核心方法
方法分两阶段：选择阶段把 2D per-head block mask 转成 per-group KV page table，执行阶段用 paged dense attention 零拷贝计算。直觉类比「图书馆取书」：block-sparse 像「按精确页号亲自取书」（并行度差），QUOKA 像「管理员挑页复印」（漏选+复印），CompactAttention 是「每组记页号清单、原地取」。技术上先用 SeerAttention 或 FlashPrefill 产生 mask $M_{b,h,i,j}$；选择阶段做 Q-block union $\bar{M}_{b,h,j}=\bigvee_i M_{b,h,i,j}$ 与 intra-group union $G_{b,g,j}=\bigvee_{h\in H(g)}\bar{M}_{b,h,j}$ 得每组 page table；执行阶段用 KV-head-major 布局 $[B, H_{kv}, L, D]$ 让每组持独立 page list，仅传 metadata 给 paged attention kernel。
核心创新是把「选 KV block」与「执行 attention」彻底解耦：把 2D block-sparse mask 仅当作 KV 选择信号而非执行计划，并通过 Q-block union、intra-group union 两个并集操作把它 lower 成 GQA-aware per-group KV page table，使选中 KV block 可被标准 paged dense attention kernel 零拷贝访问。这与已有方法的本质区别在于：block-sparse 方法直接用 mask 调度稀疏 kernel，受 $Q \ll KV$ 形态并行度限制；QUOKA 用 query 子采样 + token-level gather 绕过稀疏 kernel，但漏掉 query-specific KV 并引入 copy 开销。CompactAttention 首次实现「block-level 选择 + dense-paged 执行」组合，并通过 KV-head-major 布局让不同 group 拥有独立 page table。
方法步骤详情
方法分三步。(1) 模式打分：用 SeerAttention（SA 预训练 gate）或 FlashPrefill（FP max-threshold 算法）任一轻量 selector 对当前 chunk 的 query block 与累计 KV cache 打分，得二元 mask $M_{b,h,i,j}$。(2) 块联合构造 page table：先 Q-block union $\bar{M}_{b,h,j}=\bigvee_i M_{b,h,i,j}$，再 intra-group union $G_{b,g,j}=\bigvee_{h\in H(g)}\bar{M}_{b,h,j}$（GQA group 过大拆 4-head 子组），最终 $T_{b,g}=\{j\mid G_{b,g,j}=1\}$ 在 paged 约束下最小且保留所有原始选中块。(3) 零拷贝执行：KV cache 改 KV-head-major 布局，仅传 kv_indptr/kv_indices 给 FlashInfer paged attention kernel，原地访问 K/V payload。
技术新颖性
技术新颖性体现在三点。第一，提出 block-union 框架并证明其最小性：给定 paged 执行约束，$T_{b,g}$ 是唯一保留所有原始 2D mask 选择的最简 page table；并集牺牲稀疏度但可通过提高初始 mask 侵略性补偿（CompactAttention-FP 用 $\alpha=0.06$ FlashPrefill 初始 89.8% 稀疏，并集后 70.2%，与 FlashPrefill 自身 69.8% 相当）。第二，KV-head-major 布局首次作为「让每组持有独立 page table」的工程机制被显式提出——sequence-major 布局所有 KV head 共享同一 block table 无法表达 group-dependent 选择；用此布局把选中的 KV block 从「数据重排」变成「metadata 寻址」，消除 QUOKA 类方法的复制开销。第三，selector-agnostic 设计让方法对搜索器进步天然友好：换更便宜的 selector 只需替换第一步，page table 构造与 paged 后端无需改动。
实验结果
四点核心发现。其一，加速显著：H200 TP=2、batch=8、chunk=1024 下 CA-FP 在 128K 达 2.72× 注意力与 1.96× 端到端加速；RTX PRO 6000、batch=4、chunk=512 下 2.74×/1.51×，远超 FlashPrefill 1.32×/1.29×，随 32K→128K 单调上升（图 5）。其二，精度持平：RULER 平均 LLaMA 82.30 vs Dense 82.78、Qwen3 89.29 vs Dense 90.46；QUOKA LLaMA 平均 77.98、128K 仅 70.44（表 1）；LongBench V2 各档接近 dense（图 6）。其三，CA-FP $\alpha=0.06$ FlashPrefill 初始 89.8% 稀疏并集后 70.2%，与 FlashPrefill 69.8% 持平但加速更高。其四，执行消融（图 7c）证明同 mask 下 block-sparse 15.64 ms/layer、Copy 6.54、Metadata 5.03，in-place paged 最优。
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RULER 128K on LLaMA-3.1-8B-Instruct | RULER 准确率 (chunk=1024) | CompactAttention-FP: 74.17 | Dense: 76.59 / FlashPrefill: 74.12 / QUOKA: 70.44 | 较 FlashPrefill +0.05、较 QUOKA +3.73，与 dense 仅 2.42 分差距 |
| RULER 平均分 (32K/64K/128K) on LLaMA-3.1-8B-Instruct | RULER 平均分 | CompactAttention-FP: 82.30 | Dense: 82.78 / FlashPrefill: 82.35 / QUOKA: 77.98 | 接近 dense（-0.48），较 QUOKA +4.32 |
| RULER 平均分 on Qwen3-30B-A3B-Instruct-2507 | RULER 平均分 (32K/64K/128K) | CompactAttention-FP: 89.29 | Dense: 90.46 / FlashPrefill: 88.68 / QUOKA: 80.19 | 较 FlashPrefill +0.61、较 QUOKA +9.10 |
| 注意力加速 on H200 SXM at 128K | Attention speedup over dense | CompactAttention-FP: 2.72× | FlashPrefill: 1.32× / QUOKA: 0.61× / XAttention: 0.58× / SeerAttention: 0.44× | 较 FlashPrefill 提升 2.06×（绝对差），远超所有稀疏基线 |
| 端到端加速 on H200 SXM at 128K | End-to-end speedup over dense | CompactAttention-FP: 1.96× | FlashPrefill: 1.29× / QUOKA: 0.54× | +0.67× 绝对提升，证明 attention 层加速传导到端到端延迟 |
| LongBench V2 Hard on LLaMA-3.1-8B-Instruct | Accuracy (%) | CompactAttention-FP 接近 dense (~33) | Dense ~33 / QUOKA ~25 | Hard 子集保持精度不掉，QUOKA 明显退化 |
| 执行策略消融 (matched sparsity, RTX PRO 6000, 128K) | Layer/chunk latency (ms) | CompactAttention-FP (metadata): 5.03 ms | block-sparse kernel: 15.64 ms / CompactAttention-FP (Copy): 6.54 ms | 较 block-sparse 降低 67.8%，证明 in-place paged 执行路径最优 |
局限与改进
作者在 Section 5 承认三点局限：方法继承底层 selector 的选择质量，被 selector 漏掉的 KV block 无法通过 block-union 找回；Q-block union 与 intra-group union 用稀疏度换覆盖，并集后稀疏度会下降；有效性依赖足够大的累计 KV cache 来摊销开销，长上下文（≥32K）才能体现优势，短上下文可能因 metadata 构造劣于 dense。基于观察还可补充：方案性能强依赖 selector 质量，若初始 mask 侵略性不足，并集后 KV block 数可能接近全集丧失加速意义；GQA group 大的模型以 4-head 子组折中，若不调参或拆细可能出现 group 内 union 过度合并；论文仅在 8B 密集与 30B MoE 上验证，未覆盖 70B+ 级别或长输出生成场景；H200/Blackwell 之外的 GPU 上 paged attention 后端本身的性能上限会改变加速结论。
独立分析的弱点
独立分析三点可改进。其一，selector-agnostic 带来「上游错误放大」风险：当 SeerAttention/FlashPrefill 自身误判（如码本类或表格类长任务）时，block-union 只放大覆盖范围而不修正错误；论文仅在通用 RULER/LongBench 验证，对 SQL、代码仓库级理解等结构化长任务未做评估，建议补充 needle-in-codebase、RepoBench 等基准。其二，metadata 开销未充分表征：每 chunk 都为每个 (batch, group) 重建 ragged page list，未给出随 batch/group 数变化的扩展性曲线，若 batch×group 极大时构造开销可能反超收益。其三，KV-head-major 布局增加工程门槛：vLLM/SGLang 默认 sequence-major，直接替换会牵涉 prefix cache、beam search、多模态 KV 等子系统联动，需在 vLLM PR 中验证 prefix sharing 与 chunked prefill 共存正确性。
未来方向
结合方法特性可延伸四个方向。方向一，把 block-union 推广到 decode 阶段——当前 page table 仅在 prefill chunk 边界重建，decode 阶段若 KV cache 持续增长同样可复用 page table 抽象做 top-K 近似注意力，需要把 selector 与 paged kernel 解耦做异步更新。方向二，自适应分组大小——目前 4-head 子组是手工折中，可引入基于当前层 attention 分布熵的动态 group 划分，让并集造成的稀疏损失与 GQA group 物理结构解耦。方向三，多模态与混合模型扩展——KV-head-major 布局天然适合 cross-attention 与多模态 KV cache（每模态一个 head-major 张量即可独立构造 page table），未来可探索视觉-语言长视频理解。方向四，把 page table 构造移到 kernel 端彻底消除 CPU-GPU metadata 同步，进一步压缩小 batch 下固定开销。
复现评估
复现友好度较好。作者在 https://github.com/jiwonsong-dev/CompactAttention 公开代码，使用 FlashInfer 0.6.9 作为 paged attention 后端，硬件在 RTX PRO 6000（Blackwell SM120）与 H200 SXM（Hopper SM90）测试。LLaMA 用 SeerAttention threshold=3e-4、CompactAttention-SA threshold=5e-4、FlashPrefill α=0.01、CompactAttention-FP α=0.06；Qwen3 用 FlashPrefill α=0.02、CompactAttention-FP α=0.12；chunk size=1024 主结果、512 消融。复现难度中等偏上：需 ≥96GB 显存 GPU 跑 128K 并装 FlashInfer 0.6.9；SeerAttention gate 在 Qwen3 未公开，故仅 LLaMA 端到端可复现，Qwen3 只能用 FP 路径。门槛是硬件。
论文图表
三幅子图并列展示 (a) RULER 128K 上各方法的 accuracy–speedup 折中，CompactAttention 处于右上角最优区；(b) 在 KV=64K、稀疏度 90% 时各 block-sparse kernel 的实际提速明显低于一次性预填充与理想 10×；(c) 64K 上下文、chunk=1024 下各方法 pattern search 在 64 个 chunk 上的累计耗时。
直接给出论文要解决的三个量化问题：精度-速度折中、稀疏 kernel 在 $Q \ll KV$ 下的内核低效、pattern search 累积开销，是 motivation 的核心证据图。
(a) 用 max-attention 与 mean-attention 两种排序可视化同一窗口内 KV 位置的重要性差异，说明 query-specific 与 globally-important KV 分布不同；(b) 在 RULER Multi-key NIAH-3 与 CWE 等需要分布式信息访问的任务上，QUOKA 明显退化，而 block-sparse 方法贴近 dense。
QUOKA 类 query-subsampled 方法失败模式的实证基础，论证了「block-level + 覆盖所有 query block」设计的必要性。
整体流程图：左侧 selection 阶段把 2D per-head block mask 经过 Q-block union 与 intra-group union 转为 per-group KV block table；右侧 execution 阶段把这些 block table 作为 metadata 喂给 paged attention kernel，在原始 KV cache 上原地访问，无 explicit compaction。
方法的高层结构图，是理解「选择与执行解耦」最直观的一张图，建议放在 method 概览。
对比 sequence-major 布局 $[B, L, H_{kv}, D]$（所有 KV head 共享一个 block table，无法表达 group-dependent 选择）与 KV-head-major 布局 $[B, H_{kv}, L, D]$（每个 KV head 拥有独立连续 page，从而允许每组独立持有 page table）。
揭示为什么 KV-head-major 不是单纯实现细节，而是让独立 page table 表达得以成立的物理前提，是方法章节关键工程细节图。
在 (a) RTX PRO 6000 TP=2 batch=4 chunk=512 与 (b) H200 SXM TP=2 batch=8 chunk=1024 两套配置下，对比 dense、XAttention、SeerAttention、FlashPrefill、QUOKA、CompactAttention-SA、CompactAttention-FP 的注意力层与端到端加速；横轴为 32K/64K/128K 上下文长度。
主结果图，呈现「加速随上下文增长」这一核心卖点，并直接量化 2.72×/1.96× 等关键数字。
表格列出 Dense、QUOKA、XAttention、SeerAttention、CompactAttention-SA、FlashPrefill、CompactAttention-FP 在 LLaMA-3.1-8B-Instruct 与 Qwen3-30B-A3B-Instruct-2507 上 32K/64K/128K 及平均分；CompactAttention 两变体均接近 dense，QUOKA 显著退化。
精度的系统性证据，是论文「在加速的同时不掉点」这一论断的直接支撑，必备于 results 章节。
分组柱状图比较 Dense、QUOKA、SeerAttention、CompactAttention-SA、CompactAttention-FP 在 Overall / Easy / Hard / Short / Medium / Long 六个分组上的准确率，CompactAttention 双变体各档都接近 dense，QUOKA 在 Hard 上掉分明显。
补强 RULER 之外的第二个长上下文基准上的精度证据，说明方法不局限于 needle-in-haystack 类任务。
(a) 在 RULER 128K operating point 上 FlashPrefill 稀疏度 69.8% vs CompactAttention-FP 初始 89.8% 并集后 70.2%；(b) 在 $\alpha$ 扫描下两方法的 accuracy-attn speedup 折中曲线；(c) 在同一并集 mask 下 block-sparse kernel、CompactAttention-FP (Copy)、CompactAttention-FP (metadata) 三种执行策略的逐层耗时。
三组消融分别回答「并集后是否还有稀疏」「精确度-加速折中是否仍占优」「零拷贝执行是否真的必要」三个潜在质疑，是理解方法各组件贡献的关键。