通过共享路由索引,实现长文本推理的高效与高精度兼顾。
You Only Index Once: Cross-Layer Sparse Attention with Shared Routing
- 仅一次计算路由索引,跨层复用以降低计算开销。
- 在128K上下文下实现7.6倍解码加速与17.1倍吞吐提升。
- 适合需要长序列推理的场景,如复杂逻辑任务和文档生成。
现代大模型在长上下文推理中面临解码效率瓶颈,尤其在需生成长思维链的任务中更为显著。现有稀疏注意力方法常在效率与质量间权衡:结构化块稀疏方法虽加速明显但损失精度,而令牌稀疏方法虽准确但因全缓存top-k路由仍成本高昂。本文提出跨层稀疏注意力(CLSA),基于KV共享架构如YOCO,不仅共享键值缓存,还共享路由索引。单个索引器一次性完成令牌级top-k选择,并在各层复用,既保持令牌稀疏的细粒度选择性,又摊薄路由开销。该设计协同优化预填充、KV缓存存储及长上下文解码三大瓶颈。在短/长上下文基准测试中,CLSA兼具准确性与高效性,在128K上下文下实现最高7.6倍解码速度提升与17.1倍整体吞吐增益,为长上下文大模型提供兼顾质量与效率的完整架构方案。
原文摘要 · Abstract (English)
Long-context inference in modern LLMs is increasingly constrained by decoding efficiency, especially in reasoning-heavy settings where models generate long intermediate chains of thought. Existing sparse attention methods often face a practical efficiency-quality trade-off. Structured block sparse methods typically provide stronger acceleration but incur noticeable quality loss, while token sparse methods are usually more accurate yet deliver limited end-to-end speedup because top-k routing over the full cache remains expensive. In this work, we propose cross-layer sparse attention (CLSA), which is built on top of KV-sharing architectures such as YOCO. The core idea is to share not only the KV cache across cross-decoder layers, but also the routing index. A single indexer computes token-level top-k selection once and reuses the resulting index across layers, thereby preserving the fine-grained selectivity of token sparse attention while amortizing the routing overhead. The resulting architecture improves all major inference bottlenecks jointly, including pre-filling, KV-cache storage, and long-context decoding. Experiments across short-context and long-context benchmarks show that CLSA is both accurate and efficient, achieving up to 7.6x decoding speedup and 17.1x overall throughput improvement at 128K context. These results suggest a more complete architectural solution for long-context LLMs that jointly advances model quality and inference efficiency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。