ProxyAttn Speeds Long-Context Prefill With Shared Head Scores

Self-attention becomes expensive as prompts grow because each token can interact with every other token. Its compute and memory costs scale quadratically with sequence length, so doubling the context can roughly quadruple the work. Sparse attention reduces that cost by dividing the attention map into blocks and computing only the blocks likely to contain useful token relationships.

ProxyAttn, a training-free method from researchers at Harbin Institute of Technology and Baidu, selects those blocks using exact token-level scores from one or more representative attention heads. The selected rankings are shared with other heads, which avoids estimating importance from heavily compressed token summaries. The paper reports speedups of up to 10.3× for the attention kernel and 2.4× for prompt prefill, with little accuracy loss in its evaluations. The implementation is available on GitHub.

Block selection is the expensive guess

Sparse systems such as MInference, FlexPrefill, XAttention, and SeerAttention split attention maps into blocks, estimate each block’s importance, and compute the highest-ranked candidates. Accurate selection matters because skipped blocks cannot contribute to the output, while retaining too many blocks erodes the performance benefit.

Many selectors pool queries and keys along the sequence dimension before scoring them. Pooling is cheap, but it can average away a small number of high-scoring token pairs. The error becomes more damaging at aggressive sparsity levels, when the selector has little room to retain missed blocks. Computing every token-level dot product would provide better rankings while recreating the quadratic workload.

Heads often agree on token rankings