Prism attention speeds up long context decoding in large language models
Prism Attention: Proposal-Refined Index Sharing Mechanism for Efficient LLMs Inference
Computation and LanguageArtificial Intelligence
Summary
Large language models take a lot of time and memory to handle very long conversations or texts because they have to look back at many previous words during each step. The authors found that simply sharing which words to pay attention to across all layers causes errors because attention changes as the model processes the text. They created Prism Attention, a new way to share and refine these important word choices for each layer more efficiently. This method keeps the accuracy high while making the model run much faster, especially on very long texts.
What this means in practice
- •For llm engineers: Build large language models that decode very long texts faster by using Prism Attention's efficient attention refinement method.
- •For software engineers: Improve software using LLMs with extremely long input contexts by reducing inference time without losing accuracy.
Authors
Zhenxu Tian, Kebin Liu, Zhengwu Yang, Yi Su, Qingqing Dang, Yanlin Sha, Kaipeng Deng, Dianhai Yu, Yanjun Ma, Juntao Li, Min Zhang
Abstract
Long-context decoding in large language models (LLMs) is increasingly bottlenecked by the memory and computation required to attend over an ever-growing Key-Value (KV) cache at each autoregressive step. Although token-level sparse attention reduces this burden by restricting computation to a small subset of relevant past tokens, its practical gains are often offset by repeatedly recomputing layer-wise sparse indices as the context grows. A common remedy is to share sparse indices across layers by computing them only at designated anchor layers, but such static reuse ignores the layer-wise evolution of attention patterns, causing the reused indices to progressively deviate from the layer-specific optima and incur substantial attention loss. By revisiting cross-layer sparsity from the perspective of attention dynamics, we find that inter-layer salient token shift is highly localized: although the exact top-$κ$ indices differ across layers, the core high-attention tokens remain stable, with changes concentrated near the expanded boundary of previously selected regions. Motivated by this, we propose Prism Attention, a proposal-refined sparse attention framework. It replaces rigid cross-layer index reuse with a coarse-to-fine strategy, preserving layer-wise adaptability while retaining the efficiency benefits of cross-layer sharing. Extensive experiments show that Prism Attention achieves state-of-the-art performance among mainstream sparse attention methods on both reasoning and long-context benchmarks, while attaining a $3.4\times$ end-to-end decoding speedup over FlashAttention at 128K context length.