用压缩记忆提升长文本生成速度,降低延迟
Strong Drafts Need Compact Memories: Long-Context Speculative Decoding with Compressed KV Cache

- 为强独立草稿引入压缩键值缓存记忆,保留远距离信息
- 在32K上下文长度下,内存减少70%以上,速度提升达3.33倍
- 适合长文档摘要、多轮对话等长上下文任务
长上下文大模型应用如文档摘要和多轮对话代理需要从包含数万标记的前缀中生成内容,导致解码延迟成为主要瓶颈。推测解码(SD)可在不改变输出的情况下降低延迟,但其加速效果取决于被接受的草稿词元数量和草稿步骤延迟:轻量级草稿虽快但难以捕捉长程依赖,而强独立草稿虽能恢复接受率,却在长前缀下带来持续增长的KV访问开销。本文提出一种增强记忆的草稿机制,为强独立草稿配备压缩的草稿侧KV记忆:一个轻量适配器构建并增量更新该记忆,以保留远距离信息并精确保存近期上下文。目标验证器保留完整KV缓存,并采用标准接受/拒绝规则,保证了SD的无损性。在Llama~3.1-8B和70B模型上,前缀长度达32K的实验表明,本方法使草稿侧内存减少超过70%,相比自回归解码分别实现最高2.08倍和3.33倍的加速。
原文摘要 · Abstract (English)
Long-context LLM applications such as document summarization and multi-turn agents require generation from prefixes spanning tens of thousands of tokens, making decoding latency a major bottleneck. Speculative decoding (SD) reduces latency without changing model outputs, but its speedup depends on both accepted draft tokens and draft-step latency: Lightweight drafts are fast but lack the capacity to capture long-range dependencies, whereas strong independent drafts recover acceptance but incur growing KV-access cost at long prefixes. We introduce memory-augmented drafting for long-context SD, equipping a strong independent draft with compressed draft-side KV memory: A lightweight adaptor constructs and incrementally updates this memory to retain distant information and exact recent context. The target verifier retains its full KV cache and applies the standard accept/reject rule, preserving SD's lossless guarantee. Experiments on Llama~3.1-8B and 70B targets at prefix lengths up to 32K show that our method reduces draft-side memory by over 70%. It achieves speedups of up to 2.08x and 3.33x , respectively, over autoregressive decoding.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。