通过验证过程自动筛选关键缓存,提升长文本生成速度。
Vegas: Self-Speculative Decoding with Verification-Guided Sparse Attention
- 利用验证阶段结果动态选择关键键值对,减少计算量
- 在vLLM上提速1.25至2.81倍,优于现有方法
- 适合追求高效推理的长文本生成应用
长上下文大语言模型推理已成为当前AI应用的常态,但其性能严重受限于键值缓存(KV cache)不断增长的内存需求。已有工作表明,基于稀疏注意力的自推测解码可通过部分缓存草拟令牌并并行验证于完整缓存,实现无损加速。然而,这些方法依赖独立的缓存选择算法,忽略了验证过程中已自然计算出的每个缓存条目的重要性。本文提出Vegas,一种基于验证引导稀疏注意力的自推测解码方法。Vegas将关键缓存条目识别作为验证的副产物,并在草拟后续令牌时仅对这些条目计算注意力。这不仅提升了草稿令牌的接受率,还大幅降低缓存选择开销,从而提高解码吞吐量。Vegas在默认vLLM上实现1.25×–2.81×的吞吐量提升,在最先进的基于稀疏注意力的自推测解码方法上实现1.15×–1.29×的提升。代码已公开于https://github.com/platformxlab/vegas。
原文摘要 · Abstract (English)
Long-context large language model (LLM) inference has become the norm for today's AI applications. However, it is severely bottlenecked by the increasing memory demands of its KV cache. Previous works have shown that self-speculative decoding with sparse attention, where tokens are drafted using a subset of the KV cache and verified in parallel against the full KV cache, speeds up inference in a lossless manner. However, they rely on a standalone KV selection algorithm to select the KV entries used for drafting and overlook the fact that the criticality of each KV entry is inherently computed during verification. In this paper, we propose Vegas, a self-speculative decoding method with verification-guided sparse attention. Vegas identifies critical KV cache entries as a byproduct of verification and computes attention only over these entries when drafting subsequent tokens. This not only improves the draft token acceptance rate but also incurs low KV selection overhead, thereby improving decoding throughput. Vegas achieves a 1.25$\times$-2.81$\times$ speedup in decoding throughput over default vLLM and a 1.15$\times$-1.29$\times$ speedup over state-of-the-art sparse attention-based self-speculative decoding methods. Our code is available at https://github.com/platformxlab/vegas.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。