提出CompactAttention,提升长文本生成的注意力计算速度。
CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

- 用块联合选择机制替代传统稀疏内核,减少冗余计算
- 在128K上下文下实现最高2.72倍的注意力加速
- 适合需要高效长序列推理的部署场景
分块预填充已成为长上下文大语言模型的主流服务策略,但在此模式下的高效注意力计算仍具挑战。现有稀疏注意力方法主要针对单次预填充设计,难以适配分块预填充:块稀疏内核在查询长度受限于分块大小时效率下降,而细粒度模式搜索在每次分块累积的键值缓存上重复执行时成本高昂。QUOKA虽直接面向分块预填充,避免稀疏内核开销,但依赖查询子采样的逐标记键值选择,可能遗漏特定查询的键值项,并引入显式的键值复制开销。为此,我们提出CompactAttention,一种基于块联合键值选择的分块预填充注意力机制。CompactAttention将二维块稀疏掩码视为键值选择信号而非直接的稀疏内核执行计划,通过查询块并集与组内并集转换为适应多头注意力(GQA)的每组键值块表。该结构在分页执行约束下生成最小键值块表,保留输入掩码所选的所有键值块,支持原地访问,无需显式键值压缩。在LLaMA-3.1-8B-Instruct上,CompactAttention在RULER基准上保持接近稠密注意力的精度,同时在128K上下文长度下,分块预填充时实现高达2.72倍的注意力加速。
原文摘要 · Abstract (English)
Chunked prefill has become a widely adopted serving strategy for long-context large language models, but efficient attention computation in this regime remains challenging. Existing sparse attention methods are primarily designed for one-shot prefill and do not translate efficiently to chunked prefill: block-sparse kernels lose efficiency when the query length is limited by the chunk size, while fine-grained pattern search becomes costly when repeated over the accumulated KV cache at every chunk. QUOKA, a recent method that directly targets chunked prefill, avoids sparse-kernel overhead but relies on query-subsampled, token-level KV selection, which can miss query-specific KV entries and introduce explicit KV-copy overhead. To address these limitations, we propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection. CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans, and converts them into GQA-aware per-group KV block tables through Q-block union and intra-group union. This construction produces the minimal block tables that preserve all KV blocks selected by the input masks under paged execution constraints, enabling selected KV blocks to be accessed in place without explicit KV compaction. On LLaMA-3.1-8B-Instruct, CompactAttention maintains accuracy close to dense attention on the RULER benchmark while delivering up to 2.72$\times$ attention speedup at 128K context length under chunked prefill.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。