RRAttention: Dynamic Block Sparse Attention via Per-Head Round-Robin Shifts for Long-Context Inference

TL;DR

RRAttention introduces a dynamic sparse attention mechanism with a round-robin sampling strategy that maintains query independence, reduces computational complexity, and achieves near full attention performance on long-context tasks.

Contribution

The paper proposes RRAttention, a novel dynamic sparse attention method that combines global pattern discovery, query independence, and efficiency through a head round-robin sampling strategy.

Findings

Achieves over 99% of full attention performance.

Reduces complexity from O(L^2) to O(L^2/S^2).

Provides 2.4× speedup at 128K context length.

Abstract

The quadratic complexity of attention mechanisms poses a critical bottleneck for large language models processing long contexts. While dynamic sparse attention methods offer input-adaptive efficiency, they face fundamental trade-offs: requiring preprocessing, lacking global evaluation, violating query independence, or incurring high computational overhead. We present RRAttention, a novel dynamic sparse attention method that simultaneously achieves all desirable properties through a head \underline{r}ound-\underline{r}obin (RR) sampling strategy. By rotating query sampling positions across attention heads within each stride, RRAttention maintains query independence while enabling efficient global pattern discovery with stride-level aggregation. Our method reduces complexity from $O(L^2)$ to $O(L^2/S^2)$ and employs adaptive Top-$\tau$ selection for optimal sparsity. Extensive experiments on natural language understanding (HELMET) and multimodal video comprehension (Video-MME) demonstrate that RRAttention recovers over 99\% of full attention performance while computing only half of the attention blocks, achieving 2.4$\times$ speedup at 128K context length and outperforming existing dynamic sparse attention methods.