RBS-Attention：面向长上下文大语言模型的半径有界稀疏预填充

摘要

长上下文大语言模型推理日益受限于预填充阶段——在生成开始前，稠密自注意力需处理整个提示。稀疏块选择可降低这一开销，但块质心可能把高度相关的token隐藏在众多无关token之中。我们将这种失效模式称为均值稀释（mean dilution），并提出RBS-Attention——一种免训练的稀疏预填充方法，包含两个互补的选择分支：质心基础分支捕捉平均相关性，救援分支则利用最大key块半径及其随提示、层、头变化的分布，识别可能被低估的块。对两个分支独立设置阈值并合并其掩码，可控制救援块的贡献，同时保持常规的块稀疏FlashAttention执行。在H100 GPU上，对于128K上下文的Qwen3-30B-A3B-Instruct-2507-FP8，RBS-Attention实现了20.65倍独立预填充注意力加速、11.92倍vLLM预填充注意力加速和5.97倍端到端首token时间加速。在稠密的Qwen3-32B模型上，其RULER总体准确率为88.65（稠密注意力为89.52）；LongBench-v2、InfiniteBench和Video-MME提供了额外的质量评估。支持性实验测量了实际保留率，在相同密度下比较不同选择器，并刻画了块大小、阈值与内存行为。这些结果共同表明，半径自适应的双分支选择是长上下文预填充的有效方法。

RBS-Attention：面向长上下文大语言模型的半径有界稀疏预填充：论文配图 — (a) 不同预算下的块选择质量。(b) 仅用均值评分会低估分散的块。(c) Q5 包含 39.5% 的前 5% 块。图 1：促成 RBS-Attention 的实证观察。(a) 在给定选中块比例下，RBS 的注意力权重覆盖率高于图中展示的稀疏预填充（sparse-prefill）基线。绿色虚线表示该覆盖率诊断的参考上限，区别于式 (3) 中的逐 token logit 上界。(b) 块按半径五分位从紧凑（Q1）到分散（Q5）分组。仅用均值的排名低估随半径增大而加剧；RBS 在 Q4 和 Q5 中降低了低估程度，在 Q5 中差异达 25 个百分点。(c) 在前 5% 的注意力块中，39.5% 落在最高半径五分位，而均匀参考值为 20%。