需要多集中的注意力? Oracle 引导的稀疏预填充混合长上下文模型中的完整/GQA 层
How Much Dense Attention is Necessary? Oracle-Guided Sparse Prefill for Full/GQA Layers in Hybrid Long-Context Models
摘要
长上下文预填充仍然很昂贵,因为完整/GQA 层仍然对历史序列进行评分,即使在具有局部、稀疏、线性或循环组件的混合模型中也是如此。我们研究在明确的支持粒度和 top-k 预算下需要多少集中注意力才能保留任务级行为。我们为现有的 GQA 检查点引入了注意力质量 top-k 预言机:对于每一层和查询位置,它计算密集注意力,选择头部平均标记支持,并仅在该支持上重新计算注意力。预言机是一个诊断参考,而不是一个可部署的加速器,它将稀疏预算可行性与索引器错误和运行时实现效果分开。在 Qwen 系列检索密集型评估中,最长的每个查询预言行保持在 1 个点的密集范围内,而从 4K 到 100K 的 Qwen3.5-9B RULER 式扫描保持在 0.48 个点之内。在预言机的指导下,我们从密集的注意力质量分布中推导出由 KL 蒸馏 训练的头部塌陷辅助索引器,同时保持骨干网冻结。使用单独提取的 Qwen3.5-0.8B 和 Qwen3.5-9B 索引器,报告的 16K/32K 验证宏差距为 +2.04 和 +1.13 点,被视为质量保持而不是改进;融合选择块共享支持可能会带来更大的实现差距。初步的单卡 TTFT 测量显示,与密集的 FlashAttention-2 基线相比,NPU 上的 Qwen3.5-0.8B 的蒸馏索引器稀疏服务加速为 1.71 倍,GPU 上的 Qwen3.5-9B 的加速为 1.93 倍。额外的随机初始压力行达到 3.44 倍,表明运行时空间稀疏,但未验证输出质量。第一个版本将预言机可行性、蒸馏索引器质量和运行时空间分开,为未来的工作留下了完全匹配的质量延迟边界。