Prism Attention: Proposal-Refined Index Sharing Mechanism for Efficient LLMs Inference
作者: Zhenxu Tian, Kebin Liu, Zhengwu Yang, Yi Su, Qingqing Dang, Yanlin Sha, Kaipeng Deng, Dianhai Yu, Yanjun Ma, Juntao Li, Min Zhang
分类: cs.CL, cs.AI
发布日期: 2026-10-16
备注: Published as a main conference paper at NeurIPS 2026
💡 一句话要点
提出Prism Attention以解决长上下文解码效率问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长上下文解码 稀疏注意力 大型语言模型 计算效率 动态索引
📋 核心要点
- 现有的长上下文解码方法在处理日益增长的键值缓存时面临内存和计算的瓶颈，导致效率低下。
- 论文提出的Prism Attention框架通过粗到细的策略，优化了跨层稀疏索引的重用，增强了层级适应性。
- 实验结果显示，Prism Attention在主流稀疏注意力方法中表现最佳，并在128K上下文长度下实现了3.4倍的解码速度提升。
📝 摘要（中文）
在大型语言模型（LLMs）的长上下文解码中，内存和计算需求日益成为瓶颈。尽管基于令牌的稀疏注意力通过限制计算在相关过去令牌的子集上减轻了这一负担，但随着上下文的增长，层级稀疏索引的重复计算往往抵消了其实际收益。为了解决这一问题，本文提出了Prism Attention框架，通过粗到细的策略替代了刚性的跨层索引重用，保持了层级适应性，同时保留了跨层共享的效率优势。实验表明，Prism Attention在推理和长上下文基准测试中表现出色，相较于FlashAttention在128K上下文长度下实现了3.4倍的端到端解码加速。
🔬 方法详解
问题定义：本文旨在解决大型语言模型在长上下文解码中由于内存和计算需求导致的效率瓶颈。现有方法在处理稀疏注意力时，层级稀疏索引的重复计算使得性能提升有限。
核心思路：论文提出的Prism Attention框架通过引入粗到细的策略，避免了静态的跨层索引重用，从而保持了层级的适应性，提升了注意力计算的效率。
技术框架：该框架包括多个模块，首先在指定的锚层计算稀疏索引，然后根据上下文的变化动态调整这些索引，确保高关注令牌的稳定性和准确性。
关键创新：Prism Attention的核心创新在于其动态调整的稀疏索引机制，与现有方法的静态重用形成鲜明对比，显著减少了注意力损失。
关键设计：在设计中，论文关注于高关注令牌的局部变化，通过精细化的索引调整策略，确保了在不同层级间的有效信息共享，同时优化了计算效率。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Prism Attention在推理和长上下文基准测试中达到了最先进的性能，相较于FlashAttention在128K上下文长度下实现了3.4倍的端到端解码速度提升，显示出显著的效率优势。
🎯 应用场景
Prism Attention框架具有广泛的应用潜力，尤其在需要处理长上下文的自然语言处理任务中，如机器翻译、对话系统和文本生成等。其高效的解码能力将推动大型语言模型在实时应用中的落地，提升用户体验和系统响应速度。
📄 摘要（原文）
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.