Short context drafters retain speedup benefits at ultra long contexts
Long-Context Drafting Does Not Require Long-Context Drafters
Machine Learning
Summary
Generating very long text faster usually requires special long-memory models, which can be expensive to train. This paper shows that you can use existing short-memory text generators far beyond their usual length without changing their training. The authors introduce methods that let these short-memory drafters efficiently handle huge text contexts by cleverly adjusting internal positions and limiting what they 'pay attention' to. They prove this approach maintains speed benefits at lengths up to two million tokens across multiple datasets.
What this means in practice
- •For large language model engineers: Use short-context language models as efficient drafters for ultra-long text generation without retraining or fine-tuning.
- •For cloud ai service providers: Reduce computation costs in ultra-long context language model inference by employing this method to maximize drafting speedups.
Authors
Aref Jafari, Vansh Bhatia, Matt Qin, Subhajit Dutta Chowdhury, Ben Lee, Vikram Appia, Ali Ghodsi, Mehdi Rezagholizadeh
Abstract
Speculative decoding pairs a cheap drafter with parallel verification against the target. At ultra-long context it stops paying, because a drafter that attends over the full prefix inherits the target's context-dependent cost, so the speedup it was meant to provide first erodes and then reverses. The standard response is to build a long-context drafter, by training or fine-tuning one, or by using a multi-token-prediction head shipped inside the target checkpoint. We argue this is unnecessary, because the drafters in use are target-conditioned, reading hidden states the target computed over the entire prefix, so distant context already reaches them without their own attention, and verification corrects what they miss. We therefore freeze released short-context drafters and run them far past their training length under two inference-time mechanisms, one bounding the attention cache they see, which fixes their cost, and one re-basing their rotary position indices into the trained range. No weight is touched and the only tunable saturates early, so one vLLM-side configuration spans 4k to two million tokens, costing at most 0.16x of speedup at short length. We prove when a common index shift leaves the drafter's computation exactly unchanged, and decompose acceptance into two measurable gaps with sharp bounds. Whatever speedup a frozen drafter delivers at short context, it retains out to two million tokens, across three target-drafter pairs and five corpora. The claim concerns the flatness rather than the multiple, and removing either mechanism eliminates it. The gain is conditional on the speculative cost per emitted token staying below the matched target-only cost.