RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models
RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models
RBS-Attention：面向长上下文大语言模型的半径受限稀疏预填充
Abstract: 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.
摘要： 长上下文大语言模型的推理日益受到预填充（prefill）阶段的限制，在该阶段中，密集自注意力机制需要在生成开始前处理整个提示词。稀疏块选择（Sparse block selection）可以降低这一成本，但块质心（block centroid）可能会将高度相关的 token 掩盖在大量无关 token 之中。我们将这种失效模式称为“均值稀释”（mean dilution），并提出了 RBS-Attention，这是一种无需训练的稀疏预填充方法，包含两个互补的选择分支。
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.
一个质心基础分支用于捕捉平均相关性，而一个救援分支则利用最大键块半径（maximum key-block radius）及其在提示词、层和头维度上的分布，来识别那些容易被低估的块。通过对这两个分支进行独立的阈值处理并合并它们的掩码（mask），可以在保持常规块稀疏 FlashAttention 执行的同时，有效控制救援块的贡献。
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.
在 H100 GPU 上，RBS-Attention 在 Qwen3-30B-A3B-Instruct-2507-FP8 模型、128K 上下文长度下，实现了 20.65 倍的独立预填充注意力加速、11.92 倍的 vLLM 预填充注意力加速，以及 5.97 倍的端到端首字延迟（time-to-first-token）加速。在密集型 Qwen3-32B 模型上，其 RULER 总准确率为 88.65，而密集注意力机制为 89.52；LongBench-v2、InfiniteBench 和 Video-MME 提供了额外的质量评估。配套实验测量了实际保留率，比较了相同密度下的选择器，并表征了块大小、阈值和内存行为。综上所述，这些结果证明了半径自适应双分支选择是长上下文预填充的一种有效方法。