RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models

📄 arXiv: 2609.20971v1 📥 PDF

作者: Chuxu Song, Jiuqi Wei, Zhencan Peng

分类: cs.AI

发布日期: 2026-09-17

💡 一句话要点

提出RBS-Attention以解决长上下文大语言模型的稀疏预填充问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文 大语言模型 稀疏注意力 预填充 计算效率 自然语言处理 深度学习

📋 核心要点

  1. 现有的长上下文大语言模型在预填充阶段面临高计算成本，密集自注意力处理整个提示的效率低下。
  2. 本文提出RBS-Attention，通过质心基础分支和救援分支的双重选择机制，有效识别和利用相关性强的块，降低计算成本。
  3. 在实验中，RBS-Attention在多个基准测试中表现出显著的加速效果，尤其在H100 GPU上实现了高达20.65倍的预填充注意力加速。

📝 摘要（中文）

长上下文大语言模型的推理受到预填充的限制，密集自注意力需要处理整个提示，导致计算成本高。稀疏块选择可以降低这一成本，但块质心可能隐藏相关性强的标记。为此，本文提出RBS-Attention，一种无训练的稀疏预填充方法，包含两个互补选择分支：质心基础分支捕捉平均相关性，而救援分支利用最大关键块半径及其依赖于提示、层和头的分布来识别可能被低估的块。通过独立阈值化两个分支并结合其掩码，控制救援块的贡献，同时保留常规块稀疏FlashAttention执行。在H100 GPU上，RBS-Attention实现了20.65倍的独立预填充注意力加速和5.97倍的端到端首次令牌时间加速。

🔬 方法详解

问题定义：本文旨在解决长上下文大语言模型在预填充阶段的计算效率问题。现有方法中，密集自注意力处理整个提示，导致计算资源浪费和延迟，尤其是在处理长文本时，相关性强的标记可能被低估。

核心思路：RBS-Attention通过引入质心基础分支和救援分支的双重选择机制，旨在提高相关性标记的识别率。质心基础分支捕捉平均相关性，而救援分支则关注可能被低估的块，从而提升整体性能。

技术框架：RBS-Attention的整体架构包括两个主要模块：质心基础分支和救援分支。质心基础分支通过计算块的平均相关性来选择块，而救援分支则利用最大关键块半径及其依赖的分布来识别潜在的低估块。两个分支的掩码通过独立阈值化进行组合，以控制救援块的贡献。

关键创新：RBS-Attention的核心创新在于其双分支选择机制，能够有效识别和利用相关性强的块，避免了传统方法中因质心选择导致的平均稀释问题。这一方法在不需要额外训练的情况下，显著提高了长上下文处理的效率。

关键设计：在设计中，关键参数包括块的选择阈值、最大关键块半径以及与提示、层和头相关的分布特征。这些设计确保了RBS-Attention在保持计算效率的同时，能够有效捕捉相关性信息。实验中还测量了实际保留率，并比较了不同选择器在匹配密度下的表现。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RBS-Attention在H100 GPU上实现了20.65倍的独立预填充注意力加速，11.92倍的vLLM预填充注意力加速，以及5.97倍的端到端首次令牌时间加速。与密集Qwen3-32B模型相比，其整体RULER准确率为88.65，接近密集注意力的89.52，显示出显著的性能提升。

🎯 应用场景

RBS-Attention的研究成果在长文本处理、对话系统和自然语言生成等领域具有广泛的应用潜力。通过提高长上下文大语言模型的推理效率，该方法能够加速文本生成和理解任务，提升用户体验，尤其在实时交互场景中具有重要价值。未来，该技术可能推动更高效的语言模型开发和应用。

📄 摘要（原文）

Long-context large language model inference is increasingly limited by prefill, where dense self-attention processes the entire prompt before generation begins. Sparse block selection can reduce this cost, but a block centroid may hide a highly relevant token among many irrelevant ones. We call this failure mode mean dilution and propose RBS-Attention, a training-free sparse-prefill method with two complementary selection branches. A centroid base branch captures average relevance, while a rescue branch uses the maximum key-block radius and its prompt-, layer-, and head-dependent distribution to identify blocks at risk of underestimation. Independently thresholding the two branches and combining their masks controls the contribution of rescue blocks while preserving regular block-sparse FlashAttention execution. On H100 GPUs, RBS-Attention achieves 20.65$\times$ standalone prefill-attention speedup, 11.92$\times$ vLLM prefill-attention speedup, and 5.97$\times$ end-to-end time-to-first-token speedup at 128K on Qwen3-30B-A3B-Instruct-2507-FP8. On the dense Qwen3-32B model, it obtains 88.65 overall RULER accuracy versus 89.52 for dense attention; LongBench-v2, InfiniteBench, and Video-MME provide additional quality evaluation. Supporting experiments measure actual retention, compare selectors at matched density, and characterize block-size, threshold, and memory behavior. Together, these results support radius-adaptive dual-branch selection as an effective approach to long-context prefill.