提出稀疏预填充机制,用最小注意力计算保持长文本模型性能。
How Much Dense Attention is Necessary? Oracle-Guided Sparse Prefill for Full/GQA Layers in Hybrid Long-Context Models

- 用密集注意力质量作为参考,筛选关键上下文进行稀疏计算。
- 在10万上下文长度下,性能损失小于0.48分,接近全注意力效果。
- 适合追求长文本推理效率的开发者,尤其关注延迟优化的部署场景。
长上下文预填充仍成本高昂,因全注意力或分组查询注意力(GQA)层需对历史序列进行评分,即使在混合了局部、稀疏、线性或循环组件的模型中亦然。本文研究在明确支持粒度和前k预算下,维持任务级行为所需的密集注意力量。提出一种基于注意力质量的前kOracle:对每层和查询位置,先计算密集注意力,选取头平均的令牌支持,再仅在此支持上重新计算注意力。该方法为诊断工具,非可部署加速器,可分离稀疏预算可行性与索引误差及运行时实现的影响。在Qwen系列以检索为主的评估中,最长查询的Oracle结果与密集注意力相差不超过1分;从4K到10万上下文的Qwen3.5-9B RULER式扫描,性能下降不超过0.48分。基于此Oracle,训练一个头压缩的辅助索引器,通过知识蒸馏从密集注意力分布学习,保持主干冻结。使用分别蒸馏的Qwen3.5-0.8B与9B索引器,16K/32K验证宏得分分别提升+2.04和+1.13点,视为质量保持而非提升;融合选择块共享支持可能引入更大实现差距。初步单卡TTFT测量显示,蒸馏索引器在NPU上对Qwen3.5-0.8B实现1.71倍加速,在GPU上对Qwen3.5-9B达1.93倍,相较其密集FlashAttention-2基线。额外随机初始化压力测试达3.44倍加速,表明存在稀疏运行时潜力但未验证输出质量。本工作首次分离出Oracle可行性、蒸馏索引器质量与运行时潜力三者,留待未来实现完全匹配的质量-延迟前沿。
原文摘要 · Abstract (English)
Long-context prefill remains expensive because full/GQA layers still score the historical sequence, even in hybrid models with local, sparse, linear, or recurrent components. We study how much dense attention is needed to preserve task-level behavior under explicit support granularity and top-k budgets. We introduce an attention-mass top-k oracle for existing GQA checkpoints: for each layer and query position, it computes dense attention, selects head-averaged token support, and recomputes attention only on that support. The oracle is a diagnostic reference, not a deployable accelerator, and separates sparse-budget feasibility from indexer error and runtime realization effects. On Qwen-family retrieval-heavy evaluations, the longest per-query oracle rows stay within 1 point of dense, and a Qwen3.5-9B RULER-style sweep from 4K to 100K stays within 0.48 points. Guided by the oracle, we derive a head-collapsed auxiliary indexer trained by KL distillation from dense attention-mass distributions while keeping the backbone frozen. With separately distilled Qwen3.5-0.8B and Qwen3.5-9B indexers, the reported 16K/32K validation macro gaps are +2.04 and +1.13 points, treated as quality preservation rather than improvement; fused selection-block-shared support can introduce a larger realization gap. Preliminary single-card TTFT measurements show distilled-indexer sparse serving speedups of 1.71x for Qwen3.5-0.8B on NPU and 1.93x for Qwen3.5-9B on GPU against its dense FlashAttention-2 baseline. Additional random-init stress rows reach 3.44x, indicating sparse-runtime headroom but not validated output quality. This first release separates oracle feasibility, distilled-indexer quality, and runtime headroom, leaving a fully matched quality-latency frontier to future work.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。