Strong Drafts Need Compact Memories: Long-Context Speculative Decoding with Compressed KV Cache
2026-08-31 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address the problem of slow text generation in large language models when working with very long inputs, like long documents or many turns in a conversation. They improve a method called speculative decoding by adding a special memory component that helps the model remember important information from far back in the text while generating drafts quickly. This lets the model keep speed without losing accuracy and reduces the amount of memory needed during drafting. Their tests show their approach can generate text more than twice as fast compared to traditional methods, especially on long inputs.
large language modelsspeculative decodingdecoding latencykey-value memorylong-context generationautoregressive decodingLlama 3.1draft memoryincremental updatesaccept/reject rule
Authors
Tong Yuan, Chengxi Liao, Zeyi Wen
Abstract
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.