Research map / Large language models / Language, safety and interpretability
TokenSwift: Lossless Acceleration of Ultra Long Sequence Generation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes TokenSwift, a lossless speculative-decoding framework that accelerates generation of 20K–100K token sequences using self-drafting heads, n-gram token reutilization, dynamic partial KV cache, and a contextual repetition penalty, achieving over 3x speedup.
Problem
Generating ultra-long sequences with LLMs takes hours; existing short-sequence speculative decoding methods fail to accelerate at this length due to model reloading, growing KV cache, and repetitive generation.
Approach
Use the target model itself as the draft model with only a few trained linear heads (Medusa-style, chained residually), add n-gram retrieval from generated text for extra draft candidates, maintain a dynamically updated partial KV cache ranked by query-key importance, and apply a windowed contextual penalty to reduce repetition while keeping lossless verification.
Key results
- TokenSwift achieves up to 3.25x speedup at 100K generation on LLaMA2-7B-128k (MHA) and up to 3.38x on LLaMA3.1-8B (GQA) with prefix 8192.
- On Qwen2.5-14B, generating 100K tokens saves up to 5.54 hours relative to autoregressive decoding.
- Min-p sampling with contextual penalty raises average Distinct-n from 0.12 to 0.69 at about 8% speedup cost.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add token reutilization (k=20) rather than disable it (k=0) (ablation) | TokenSwift · LLaMA3.1-8B, ultra-long generation | speedup | Helped |
| use dynamic partial KV cache updates instead of a one-time partial cache update (ablation) | TokenSwift · LLaMA3.1-8B, 100K tokens | speedup: 1.44 -> 1.96 | Helped |
| use dynamic partial KV cache updates instead of retaining the full cache (ablation) | TokenSwift · LLaMA3.1-8B, 100K tokens | speedup: 1.47 -> 1.96 | Helped |
| add contextual penalty (θ=1.2, W=1024) (ablation) | TokenSwift with min-p sampling · LLaMA3.1-8B, 100K tokens | average Distinct-n: 0.12 -> 0.69 | Helped |
| add contextual penalty (θ=1.2, W=1024) (ablation) | TokenSwift with min-p sampling · LLaMA3.1-8B, 100K tokens | speedup: 3.58 -> 3.27 | Hurt |
| use tree [1,3,3,3] instead of [1,9,9,9] (ablation) | TokenSwift · LLaMA3.1-8B, 100K tokens, k=0 | speedup: 0.96 -> 1.96 | Helped |
| use tree [1,3,3,3] instead of [3,3,3,3] (ablation) | TokenSwift · LLaMA3.1-8B, 100K tokens, k=0 | speedup: 1.91 -> 1.96 | Helped |
| vary temperature across 0.4–1.2 (ablation) | TokenSwift with top-p sampling · LLaMA3.1-8B, 100K tokens | acceptance rate: 0.89, 0.90, 0.85, 0.87, 0.83 | Mixed |
| increase prefill length from 2048 to 8192 while disabling token reutilization (ablation) | TokenSwift · LLaMA3.1-8B, 100K tokens | speedup: 1.87 -> 1.77 | Hurt |
| increase penalty window from W=20 to W=2048 (ablation) | TokenSwift · LLaMA3.1-8B, 100K tokens | speedup: 3.58 -> 3.31 | Hurt |
| increase penalty window from W=20 to W=2048 (ablation) | TokenSwift · LLaMA3.1-8B, generated sequence diversity | average Distinct-n: 0.79 -> 0.80 | Mixed |
Surprising findings
- TriForce* achieves only about 1.0–1.25x speedup on LLaMA3.1-8B and below 1.0x at some lengths, and suffers low quality without repetition penalty.
- Medusa* gives only ~0.85–1.45x speedup, showing standard Medusa heads do not translate to ultra-long generation.
Limitations
- Speedup depends on acceptance rate, which drops for sampling with higher temperature.
- Repetition is not eliminated; contextual penalty only mitigates it, and penalties reduce acceptance or speed in some settings.
- Evaluation is limited to PG-19 and base (non-instruct) models.
Experimental setup
Datasets: PG-19 Wikipedia (20231101.en) C4-en
Benchmarks: PG-19
Metrics: acceptance rate speedup Distinct-n
Baselines: TriForce* Medusa* autoregressive decoding
Builds on: Medusa TriForce Speculative decoding SnapKV
Closest papers by meaning
- Gumiho: A Hybrid Architecture to Prioritize Early Tokens in Speculative Decoding (ICML 2025 · similarity 0.90)
- MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding (ICLR 2025 · similarity 0.90)
- SPEED-Bench: A Unified and Diverse Benchmark for Speculative Decoding (ICML 2026 · similarity 0.89)
- Traversal Verification for Speculative Tree Decoding (NeurIPS 2025 · similarity 0.89)
- Optimized Multi-Token Joint Decoding With Auxiliary Model for LLM Inference (ICLR 2025 · similarity 0.88)
- Judge Decoding: Faster Speculative Sampling Requires Going Beyond Model Alignment (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.