RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models
Organizations: Department of Computer Science, Rutgers University · OceanBase, Ant Group
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. 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 standalone prefill-attention speedup, 11.92 vLLM prefill-attention speedup, and 5.97 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.
Figures & tables
| Target | FlashPrefill | XAttn | FlexPrefill | MInference | RBS-Attention |
|---|---|---|---|---|---|
| 1% | 1.77 / 56.65 | 1.02 / 44.51 | 1.17 / 31.64 | 1.01 / 50.83 | 1.80 / 60.36 |
| 2% | 1.91 / 63.32 | 1.71 / 50.28 | 2.25 / 41.41 | 1.82 / 68.58 | 2.03 / 70.51 |
| 5% | 4.69 / 75.36 | 3.81 / 60.47 | 3.59 / 50.72 | 5.90 / 76.41 | 5.34 / 80.28 |
| 10% | 10.19 / 85.01 | 7.63 / 71.01 | 6.69 / 65.87 | 11.45 / 87.18 | 11.03 / 88.36 |
| Accuracy | Estimated density (%) | ||||||
|---|---|---|---|---|---|---|---|
| Method | Short | Medium | Long | Total | Mean | p50 | p90 |
| Qwen3-32B | |||||||
| Dense | 0.450 | 0.367 | 0.380 | 0.394 | 100 | 100 | 100 |
| FlashPrefill | 0.378 | 0.349 | 0.361 | 0.362 | 14.149 | 9.197 | 28.047 |
| XAttn | 0.439 | 0.340 | 0.370 | 0.376 | 29.197 | 26.711 | 43.858 |
| FlexPrefill | 0.422 | 0.293 | 0.352 | 0.352 | 11.397 | 10.203 | 16.260 |
| Method | 4K | 8K | 16K | 32K | 64K | 128K | Overall |
|---|---|---|---|---|---|---|---|
| Dense | 93.30 | 91.90 | 92.00 | 89.70 | 89.50 | 80.70 | 89.52 |
| density | 100 | 100 | 100 | 100 | 100 | 100 | 100 |
| FlashPrefill | 91.00 | 89.20 | 88.80 | 81.50 | 81.80 | 68.10 | 83.40 |
| density | 81.091 | 59.781 | 40.970 | 26.754 | 16.872 | 10.075 | 39.257 |
| XAttn | 92.40 | 91.30 | 90.90 | 86.80 | 87.20 | 79.20 | 87.97 |
| density | 69.518 | 59.482 | 51.776 | 43.645 | 36.201 | 29.196 | 48.303 |
| Task | Dense | FlashPrefill | XAttn | RBS-Attention |
|---|---|---|---|---|
| Ret.KV | 0.244 | 0.148 | 0.194 | 0.204 |
| Ret.PassKey | 1.000 | 1.000 | 1.000 | 1.000 |
| Ret.Num | 1.000 | 1.000 | 1.000 | 0.998 |
| En.Sum | 0.289 | 0.291 | 0.272 | 0.296 |
| En.MC | 0.105 | 0.100 | 0.105 | 0.092 |
| En.QA | 0.087 | 0.075 | 0.084 | 0.083 |
| Method | Short | Medium | Long | Total |
|---|---|---|---|---|
| Dense | 73.6 | 62.5 | 58.6 | 66.16 |
| FlashPrefill | 71.2 | 61.5 | 60.9 | 65.26 |
| XAttn | 72.8 | 61.0 | 61.2 | 65.76 |
| FlexPrefill | 71.0 | 57.9 | 56.3 | 62.82 |
| MInference | 71.4 | 60.5 | 59.8 | 64.72 |
| RBS-Attention | 72.7 | 61.2 | 62.1 | 65.98 |
| Selector | Actual density | Accuracy | 95% CI |
|---|---|---|---|
| Centroid-only | 5.300 | 73.63 | [68.00, 81.21] |
| Full-L2 bound | 5.295 | 72.41 | [66.32, 79.79] |
| Quest-style | 5.334 | 73.71 | [66.76, 80.47] |
| RBS adaptive union | 5.350 | 76.67 | [69.01, 82.15] |
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
| Method | 2.03% anchor | 5.34% anchor | 11.03% anchor |
|---|---|---|---|
| FlashPrefill | 2.038 / 67.13 | 5.300 / 74.73 | 11.035 / 79.62 |
| MInference | 1.944 / 68.40 | 5.286 / 76.87 | 10.950 / 83.44 |
| XAttn | — | 5.311 / 66.77 | — |
| FlexPrefill | — | 5.352 / 58.56 | — |
| RBS-Attention | 2.030 / 70.51 | 5.340 / 80.28 | 11.030 / 88.36 |
| Selector | Retrieval accuracy | Multi-key accuracy |
|---|---|---|
| Centroid-only | 88.75 | 90.00 |
| Full-L2 bound | 87.81 | 90.00 |
| Quest-style | 88.75 | 90.00 |
| RBS adaptive union | 90.00 | 93.33 |
| Method | Target | Sys./acc. density | Selector | Sparse attn. | TTFT gain | Accuracy |
|---|---|---|---|---|---|---|
| Dense | 100 | 100 / 100 | — | 100.00 | 1.00 | 86.31 |
| FlashPrefill | 2 | 1.826 / 1.910 | 23.87 | 76.13 | 4.89 | 63.32 |
| 5 | 5.071 / 4.690 | 12.39 | 87.61 | 2.90 | 75.36 | |
| 10 | 11.147 / 10.190 | 5.99 | 94.01 | 2.04 | 85.01 | |
| MInference | 2 | 1.568 / 2.022 | 41.11 | 58.89 | 4.84 | 68.54 |
| 5 | 4.184 / 5.056 | 26.61 | 73.39 | 3.58 | 76.09 |
| Context | Batch | Dense wall time (ms) | Total prompt tokens | Dense tokens/s |
|---|---|---|---|---|
| 32K | 1 | 1,920.68 | 31,286 | 16,289 |
| 64K | 1 | 5,734.91 | 62,486 | 10,896 |
| 128K | 1 | 19,709.12 | 125,486 | 6,367 |
| 32K | 4 | 6,996.10 | 118,097 | 16,880 |
| 64K | 4 | 21,265.25 | 239,377 | 11,257 |
| 128K | 4 | 72,127.35 | 478,727 | 6,637 |
| Selector | Block size | Calibration parameter | Actual density | Accuracy |
|---|---|---|---|---|
| Centroid-only | 64 | 0.100 | 5.302 | 74.17 |
| 128 | 0.194 | 5.300 | 74.73 | |
| 256 | 0.400 | 5.276 | 72.15 | |
| RBS union | 64 | 0.096 | 5.343 | 71.86 |
| 128 | Reference | 5.340 | 80.28 | |
| 256 | 0.382 | 5.361 | 70.73 |
| Model | Setting | Mean density | p90 density | Accuracy | ||
|---|---|---|---|---|---|---|
| Qwen3-30B-A3B | Conservative | 0.16 | 0.24 | 6.478 | 7.282 | 75.31 |
| FP8 | Default | 0.22 | 0.18 | 4.997 | 5.441 | 71.86 |
| Aggressive | 0.24 | 0.16 | 4.676 | 5.048 | 75.12 | |
| Qwen3-32B | Conservative | 0.16 | 0.24 | 7.983 | 8.845 | 65.05 |
| Default | 0.22 | 0.18 | 6.029 | 6.561 | 63.90 | |
| Aggressive | 0.24 | 0.16 | 5.600 | 6.078 | 61.74 |
| Method | Configurations | Wall hours | GPU-hours | Target density | Max. error (pp) |
|---|---|---|---|---|---|
| RBS-Attention | 3 | 0.443 | 0.885 | 5% | 0.0078 |
| FlashPrefill | 3 | 0.418 | 0.837 | 5% | 0.1625 |
| MInference | 3 | 0.232 | 0.464 | 2/5/10% | 0.0913 |
| Selector | Total selector time (ms) |
|---|---|
| FlashPrefill | 349.89 |
| Full-L2 | 279.21 |
| RBS-Attention | 694.81 |
| Quest-style | 1,195.88 |
| Selector | Statistics per block | MiB/layer | Workspace (MiB) |
|---|---|---|---|
| Centroid-only | 128-vector | 1.000 | 515.499 |
| Full-L2 | 128-vector + radius | 1.016 | 515.324 |
| RBS-Attention | 128-vector + radius + | 1.031 | 1,155.601 |
| Quest-style | Two 128-vectors | 2.000 | 428.231 |
| Method | 32K | 64K | 128K |
|---|---|---|---|
| Dense | 55.710 / 56.650 | 57.328 / 59.279 | 61.552 / 65.510 |
| Centroid-only | 55.248 / 57.010 | 56.865 / 58.918 | 61.090 / 63.441 |
| RBS-Attention | 54.644 / 56.357 | 56.262 / 58.266 | 60.486 / 64.016 |
| Quest-style | 59.894 / 61.541 | 61.512 / 63.449 | 61.330 / 63.496 |