arXiv:2502.17421cs.CLcs.AI2025-02ACL被引 17

解决长文本推理中推测解码的内存与效率难题,提速超3倍。

LongSpec: Long-Context Lossless Speculative Decoding with Efficient Drafting and Verification

  • 用固定大小缓存降低草稿模型内存开销
  • 在五个数据集上实现最高3.26倍加速
  • 适合需要长上下文推理的智能体应用

随着大语言模型(LLMs)处理超长上下文的能力提升,高效推理变得愈发重要,尤其在依赖此能力的LLM智能体等新兴应用中。推测解码(SD)相比量化和模型级联等有损方法,提供了无损加速的前景。然而,现有主流SD方法大多在短文本(通常少于4k token)上训练,难以适配长上下文场景。具体挑战包括:(1) 草稿模型因大键值(KV)缓存导致内存消耗过高;(2) 短文本训练与长文本推理之间的不匹配导致性能下降;(3) 树注意力机制在处理长序列时效率低下。本文提出LongSpec框架,通过三项核心创新应对:(1) 基于常量大小KV缓存的内存高效草稿模型;(2) 新设计的位置索引缓解训练-推理不匹配问题;(3) 结合快速前缀计算与标准树注意力的注意力聚合策略,实现高效解码。实验表明,LongSpec在五个长上下文理解数据集上相比强基线Flash Attention最高提速3.26倍,并在AIME24长推理任务上使QwQ模型的墙钟时间减少2.25倍,显著提升长上下文应用的延迟表现。代码已开源:https://github.com/sail-sg/LongSpec。

原文摘要 · Abstract (English)

As Large Language Models (LLMs) can now process extremely long contexts, efficient inference over these extended inputs has become increasingly important, especially for emerging applications like LLM agents that highly depend on this capability. Speculative decoding (SD) offers a promising lossless acceleration technique compared to lossy alternatives such as quantization and model cascades. However, most state-of-the-art SD methods are trained on short texts (typically fewer than 4k tokens), making them unsuitable for long-context scenarios. Specifically, adapting these methods to long contexts presents three key challenges: (1) the excessive memory demands posed by draft models due to large Key-Value (KV) cache; (2) performance degradation resulting from the mismatch between short-context training and long-context inference; and (3) inefficiencies in tree attention mechanisms when managing long token sequences. This work introduces LongSpec, a framework that addresses these challenges through three core innovations: a memory-efficient draft model with a constant-sized KV cache; novel position indices that mitigate the training-inference mismatch; and an attention aggregation strategy that combines fast prefix computation with standard tree attention to enable efficient decoding. Experimental results confirm the effectiveness of LongSpec, achieving up to a 3.26x speedup over strong Flash Attention baselines across five long-context understanding datasets, as well as a 2.25x reduction in wall-clock time on the AIME24 long reasoning task with the QwQ model, demonstrating significant latency improvements for long-context applications. The code is available at https://github.com/sail-sg/LongSpec.

推理加速长文本推测解码KV缓存

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。