技术进展

RBS-Attention:长上下文稀疏预填充新方法

Heooo 09月21日12时02分 81 阅读

「arXiv 发布 RBS-Attention 研究,提出免训练的双分支稀疏预填充方法,针对块质心掩盖关键 token 的均值稀释问题,在 H100 上实现最高 20.65 倍预填充注意力加速,128K 上下文端到端首 token 延迟加速 5.97 倍,并在多项长上下文基准上保持高质量。」

长上下文大语言模型推理正面临一个日益突出的瓶颈:预填充阶段。在生成第一个 token 之前,稠密自注意力需要处理完整提示,序列越长,这部分开销就越难以承受。围绕这一问题,arXiv 上的一项新研究提出了 RBS-Attention,尝试用半径约束的稀疏预填充策略来削减预填充成本。

稀疏块选择是常见的降本思路,但研究者指出了它的一个典型失效模式:块质心可能掩盖大量无关 token 中的极少数高度相关 token。换句话说,当一个注意力块内部只有个别 token 真正值得关注时,基于质心计算的块级相关性评分会被平均掉,导致这个块被错误地判定为不重要。研究团队将这一现象命名为均值稀释(mean dilution)。

针对均值稀释,RBS-Attention 设计了两条互补的选择分支。第一条是质心基分支,负责捕捉平均相关性,延续了传统块稀疏选择的思路。第二条是救援分支,它利用最大键块半径,以及该半径随提示、层、头变化的分布,来识别那些存在被低估风险的块。两条分支分别独立设置阈值,再将各自的掩码组合起来,以此控制救援块对最终选择结果的贡献比例,同时保留常规的块稀疏 FlashAttention 执行路径。这种设计的关键在于免训练:不需要额外的微调或校准训练过程,即可直接作用于现有模型。

在硬件加速效果方面,论文给出的数据相当可观。在 H100 GPU 上,针对 Qwen3-30B-A3B-Instruct-2507-FP8 模型,RBS-Attention 实现了 20.65 倍的独立预填充注意力加速,以及在 vLLM 框架下 11.92 倍的预填充注意力加速。在 128K 上下文长度下,端到端的首 token 时间(time-to-first-token)加速达到 5.97 倍。首 token 延迟是长上下文交互体验中最直观的指标之一,这一数字意味着超长提示的响应等待时间可以被压缩到原来的六分之一左右。

速度提升之外,质量保持同样值得关注。在稠密的 Qwen3-32B 模型上,RBS-Attention 取得了 88.65 的 RULER 整体准确率,而稠密注意力为 89.52,两者差距不到一个百分点。研究还通过 LongBench-v2、InfiniteBench 和 Video-MME 等基准提供了额外质量评估,覆盖长文本理解、超长上下文以及视频多模态等场景。这说明稀疏预填充并非必须以明显的精度损失为代价。

论文的支撑实验进一步刻画了方法的内部机制与工程特性,包括测量实际保留率、在匹配密度条件下比较不同选择器,以及分析块大小、阈值与内存行为的影响。这些实验从多个角度验证了半径自适应双分支选择在长上下文预填充中的有效性。

从技术路线来看,RBS-Attention 的价值在于把问题定位到了块稀疏方法的评分机制本身。以往的稀疏注意力大多聚焦于选择哪些块,而这项研究指出,块的表示方式才是被低估的环节:均值统计会系统性遗漏尖峰相关性,而最大半径提供了一种补充信号。两条分支独立阈值化再合并的做法,也让稀疏度控制变得更加灵活,可以在不同层、不同头上做出差异化决策。对于需要处理长文档、长代码库或长视频序列的应用而言,这种免训练、可直接部署的预填充优化方案,提供了一条兼顾吞吐与质量的现实路径。

# 稀疏注意力 # 长上下文 # 推理加速

来源:Heooo AI工具导航