动态选择关键记忆项,让大模型长文本推理快23倍
TokenSelect: Efficient Long-Context Inference and Length Extrapolation for LLMs via Dynamic Token-Level KV Cache Selection
- 按注意力重要性动态筛选关键记忆项,不需训练
- 推理速度提升23.84倍,端到端延迟快2.28倍
- 适合需要长文本处理的AI应用开发者
大型语言模型在实际应用中对长上下文序列的需求日益增长,但面临两大挑战:超出分布的序列长度导致性能下降,以及注意力计算的二次复杂度引发过长推理时间。本文提出无需训练的动态令牌级键值缓存选择方法(TokenSelect),基于非连续注意力稀疏性,利用查询与键的点积衡量每个头的键值缓存重要性,通过每头软投票机制仅选择少数关键缓存项参与计算,保持精度。为加速选择过程,设计了基于连续查询相似性的选择缓存,并实现高效的分页点积核,显著降低选择开销。全面评估表明,TokenSelect在注意力计算上最高提速23.84倍,端到端延迟最快提升2.28倍,性能优于现有先进方法。
原文摘要 · Abstract (English)
Rapid advances in Large Language Models (LLMs) have spurred demand for processing extended context sequences in contemporary applications. However, this progress faces two challenges: performance degradation due to sequence lengths out-of-distribution, and excessively long inference times caused by the quadratic computational complexity of attention. These issues limit LLMs in long-context scenarios. In this paper, we propose Dynamic Token-Level KV Cache Selection (TokenSelect), a training-free method for efficient and accurate long-context inference. TokenSelect builds upon the observation of non-contiguous attention sparsity, using QK dot products to measure per-head KV Cache criticality at token-level. By per-head soft voting mechanism, TokenSelect selectively involves a few critical KV cache tokens in attention calculation without sacrificing accuracy. To further accelerate TokenSelect, we design the Selection Cache based on observations of consecutive Query similarity and implemented the efficient Paged Dot Product Kernel, significantly reducing the selection overhead. A comprehensive evaluation of TokenSelect demonstrates up to $23.84\times$ speedup in attention computation and up to $2.28\times$ acceleration in end-to-end latency, while providing superior performance compared to state-of-the-art long-context inference methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。