PaperlayerView exact revision on arXiv

Source-linked AI summary

TokenSwift: Lossless Acceleration of Ultra Long Sequence Generation

Tong Wu, Junzhe Shen, Zixia Jia, Yuxuan Wang, Zilong Zheng

arXiv:2502.18890v2cs.CL

TL;DR

Generating ultra-long sequences with LLMs is time-intensive, and existing speculative decoding methods are not designed to address the 100K-token setting. TOKENSWIFT combines multi-token generation, dynamic KV-cache updates, and contextual penalties, achieving over 3x speedup across model scales and architectures.

  • Problem

    Generating sequences up to 100K tokens is highly time-intensive, while existing speculative decoding methods are generally tailored to shorter sequences and face KV-cache budget constraints at ultra-long lengths.

    Source: §1 · p. 2

  • Method

    TOKENSWIFT uses multi-token generation and token reutilization, dynamic partial-KV-cache updates, and contextual penalties for ultra-long sequence generation.

    Source: §1 · p. 2

  • Results

    Over 3x acceleration across model scales and architectures, with up to 5.54 hours saved for a 14B model generating 100K tokens.

    Source: §4.2 · p. 8, §6 · p. 13

  • Takeaways & Limitations

    TOKENSWIFT is presented as a scalable solution for ultra-long generation that maintains target-model quality and mitigates repetitive content.

    Source: §1 · p. 2, §6 · p. 13

    TOKENSWIFT aims to mitigate repetition rather than eliminate it, since repetition is treated as inherent to LLMs and outside the lossless-acceleration objective.

    Source: §2 · p. 3

Abstract

from arXiv · show

Generating ultra-long sequences with large language models (LLMs) has become increasingly crucial but remains a highly time-intensive task, particularly for sequences up to 100K tokens. While traditional speculative decoding methods exist, simply extending their generation limits fails to accelerate the process and can be detrimental. Through an in-depth analysis, we identify three major challenges hindering efficient generation: frequent model reloading, dynamic key-value (KV) management and repetitive generation. To address these issues, we introduce TOKENSWIFT, a novel framework designed to substantially accelerate the generation process of ultra-long sequences while maintaining the target model's inherent quality. Experimental results demonstrate that TOKENSWIFT achieves over 3 times speedup across models of varying scales (1.5B, 7B, 8B, 14B) and architectures (MHA, GQA). This acceleration translates to hours of time savings for ultra-long sequence generation, establishing TOKENSWIFT as a scalable and effective solution at unprecedented lengths. Code can be found at https://github.com/bigai-nlco/TokenSwift.

1. Introduction

Generating ultra-long sequences with LLMs is time-intensive, and existing speculative decoding approaches face limits at these lengths. TOKENSWIFT addresses the bottleneck with methods targeting model reloading, KV-cache growth, and repetitive generation.

1. Introduction

  • Generating 100K tokens with LLaMA3.1-8B can take approximately five hours, making ultra-long generation time-intensive for practical applications.The introduction frames reducing this generation time as important for real-world LLM applications.Source: §1 · p. 2
  • Prior speculative decoding methods target shorter sequences; extending them to 100K tokens encounters KV-cache budget constraints, motivating lossless, model-agnostic acceleration with minimal training overhead.The introduction also notes only marginal acceleration gains for short-sequence generation with optimized KV-cache architectures such as GQA.Source: §1 · p. 2
  • TOKENSWIFT combines multi-token generation and token reutilization, dynamic partial-KV updates, and contextual penalties to address reloading, KV-cache loading, and repetitive outputs.Its stated design aims to draft multiple tokens in one forward pass while mitigating repetition and preserving output diversity.Source: §1 · p. 2
  • Over 3x speedup across model scales and architectures, with 100K-token LLaMA3.1-8b generation reduced from nearly five hours to 90 minutes.TOKENSWIFT also reports increasing speedup with generation length and enhanced diversity.Source: §1 · p. 2

2. Challenges

The section identifies three obstacles to ultra-long generation: model-weight reloads, a growing KV cache, and repetitive outputs. It also limits the repetition objective to mitigation rather than elimination.

2. Challenges

  • Each autoregressive token requires loading the full model from GPU storage to the computing unit, so I/O memory access—not computation—is the primary generation bottleneck.For a 100K-token sequence, the GPU must reload model weights over 100,000 times.Source: §2 · pp. 2–3
  • A one-time KV-cache compression strategy can support long prefixes and short outputs, but growing cache size exceeds its allocated budget for ultra-long outputs.Efficient generation therefore requires deciding when and how to update the partial KV cache dynamically within a limited budget.Source: §2 · p. 3
  • At considerable lengths such as 100K tokens, autoregressive generation tends toward repetitive sentences and can become bland, incoherent, or stuck in loops.The section describes repetition as an inherent LLM problem that this work aims to mitigate, not eliminate.Source: §2 · p. 3

3. TOKENSWIFT

TOKENSWIFT integrates tailored solutions for ultra-long generation into a single framework. It drafts tokens with the target model and validates them against that model’s predictions to preserve lossless output.

3.1. Overview

  • Using the target LLM with a partial KV cache for drafting and a full cache for exact-match verification makes the framework’s output align with target-model predictions.After verification, the framework selects a longest valid draft and updates the n-gram table and KV cache.Source: §3.1 · pp. 3–4, Figure 2 · p. 4

3.2. Multi-token Generation and Token Reutilization

  • Residual linear layers let the target LLM draft multiple tokens per forward pass, while frequent n-gram reuse supplies additional draft candidates.For long-prefix tasks, limited decoding steps reduce n-gram acceptance opportunities; the method is described as particularly suitable for ultra-long generation.Source: §3.2 · pp. 4–5

3.3. Dynamic KV Cache Management

  • A fixed KV-cache budget preserves the initial |S| pairs and replaces less-important pairs as generation proceeds.Importance is based on query-key dot products; for GQA, queries are partitioned rather than replicating keys, with their scores aggregated.Source: §3.3 · p. 5

3.4. Contextual Penalty and Random N-gram Selection

  • A penalty on the most recent W tokens aims to mitigate repetition while maintaining diversity, and random selection among multiple valid n-grams supports diverse, accurate output.The penalty approach modifies the all-generated-token penalty described as limiting appropriate sampling for ultra-long sequences.Source: §3.4 · pp. 5–6

4. Experiments

Experiments evaluate TOKENSWIFT across models, sequence lengths, and baselines, then examine how token reutilization, KV updates, contextual penalty, and decoding choices affect speed and output quality.

4.1. Setup

  • Experiments use PG-19 Base models, train three extra decoding heads on up to 8K tokens, and evaluate acceptance, speedup, and diversity on one A100.The comparisons include an improved TriForce baseline and a retrained Medusa baseline.Source: §4.1 · pp. 6–7

4.2. Main Results

  • Speedup grows with generation length as TOKENSWIFT mitigates increasing KV-cache loading time and n-gram acceptance improves.The main-results evaluation averages five experiments and reports standard deviations.Source: §4.2 · pp. 7–8, Figure 3 · p. 8
  • Up to 5.54 hours are saved for the 14B model; TOKENSWIFT outperforms baselines across generation lengths and model architectures, with virtually no impact from varying prefix lengths.Table 3 covers LLaMA2 and LLaMA3.1 across prefix lengths and 20K–100K outputs; Table 4 reports Qwen2.5 results at different scales.Source: §4.2 · pp. 7–8, Table 3 · p. 7, Table 4 · p. 8

4.3.1. TOKEN REUTILIZATION

  • Removing token reutilization (k = 0) reduces acceptance and speedup, while TOKENSWIFT at k = 20 shows increasing acceptance as sequences grow.The reported ablation links token reutilization to reducing model-reloading frequency.Source: Figure 3 · p. 8, §4.3.1 · p. 9

4.3.2. DYNAMIC KV UPDATES

  • Dynamic Partial Cache balances acceptance and speedup, unlike Partial Cache’s low acceptance or Full Cache’s overhead.The comparison evaluates full, partial, and dynamically updated partial KV caches during drafting.Source: §4.3.2 · p. 9

4.3.3. CONTEXTUAL PENALTY

  • Contextual penalty improves diversity across sampling methods; with min-p, Distinct-n rises from 0.12 to 0.69 at an 8% speedup compromise.The penalty is presented as orthogonal to min-p, top-p, and η-sampling; the paper also notes that greater diversity is associated with lower speedup.Source: §4.3.3 · p. 9, Table 6 · p. 9

4.4.1. TREE CONFIGURATION

  • The authors choose [1,3,3,3] to balance acceptance and verification efficiency; [1,9,9,9] has the highest acceptance but lowest speedup.They use a simple 3-ary tree because finding Medusa’s optimal tree is time-consuming and has limited impact on ultra-long generation acceleration.Source: §4.4.1 · p. 10

4.4.2. N-GRAM CANDIDATES

  • Higher k raises n-gram acceptance but can burden verification; k = 5 can exceed k = 20 in acceptance and speedup while increasing repetition.With W = 1024, raising θ toward 1.2 improves diversity; in the case study, repetition appears at about 60K tokens with the penalty versus about 5K without it.Source: Table 8 · p. 10, §4.4.2 · p. 11, §4.4.3 · p. 11

5. Related Works

The section surveys speculative-decoding and long-sequence methods, then highlights why existing KV-cache approaches do not fit tasks with both extended inputs and lengthy outputs. It also reviews planning and sampling approaches relevant to long-form generation and repetition mitigation.

5.1. Speculative Decoding

  • Speculative-decoding work spans draft-model verification, self-speculation, parallel decoding, retrieval-augmented drafting, training-centric methods, and system-level optimizations.Examples include hierarchical KV caching in Triforce (Sun et al., 2024a) and phrase retrieval in NEST (Li et al., 2024b).Source: §5.1 · p. 11
  • Triforce (Sun et al., 2024a) and MagicDec target long prefixes and short outputs, not tasks requiring both extended contexts and lengthy generation.The section identifies this mismatch as a limitation for ultra-long sequence generation.Source: §5.1 · p. 12

5.2. Long Sequence Generation

  • LongWriter (Bai et al., 2024) uses task decomposition for texts exceeding 20,000 words, while Temp-Lora (Wang et al., 2024b) adapts parameters through inference-time temporary LoRA modules.Source: §5.2 · p. 12
  • PLANET (Hu et al., 2022) uses dynamic content planning with sentence-level bag-of-words objectives to improve coherence in opinion articles and argumentative essays.Source: §5.2 · p. 13
  • Sampling approaches address repetition by truncating candidate tokens, with variants including entropy-guided η-sampling, Min-p thresholds, and Mirostat controllers.The section frames its own analysis around the balance between computational overhead and repetition suppression in ultra-long generation.Source: §5.2 · p. 13

6. Conclusion

The conclusion presents TOKENSWIFT as a lossless acceleration framework for ultra-long LLM generation, reporting over 3× acceleration across model scales and architectures. The appendix explains the lossless sampling property of speculative decoding through draft-token acceptance and rejection cases.

6. Conclusion

  • TOKENSWIFT is presented as a lossless acceleration framework for ultra-long LLM generation, with experiments reporting over 3× acceleration across model scales and architectures.Source: §6 · p. 13
  • TOKENSWIFT is reported to mitigate repetitive content while maintaining the quality and coherence of generated sequences.Source: §6 · p. 13

A. Lossless Nature of Speculative Decoding

  • The appendix states that speculative decoding preserves the target model’s sampling distribution by combining draft-token acceptance and rejection cases into an overall probability.Source: A. Lossless Nature of Speculative… · p. 18

B.1. Training Details

Training keeps the LLM parameters fixed while fine-tuning three linear layers. Inference uses 4-grams, a configurable retrieval limit, and a ternary full tree for tree attention.

B.1. Training Details

  • Training fine-tunes only three linear layers while keeping the LLM parameters fixed, using an NVIDIA A100-SXM4-80GB GPU.The hyperparameters are listed in Table 9 and are stated not to require extensive tuning.Source: §B.1 · p. 18, Table 9 · p. 18

B.2. Inference Details

  • Inference uses 4-grams for consistency with multi-token generation, and k denotes the maximum number of retrieved n-grams for token reutilization.Source: §B.2 · p. 19, Table 10 · p. 19
  • Tree attention uses a simple ternary full-tree configuration.Source: §B.2 · p. 19

C. Pre-training Details of the Llama3.1 Draft Model

The TriForce draft model used for LLaMA3.1-8b is a 250M-parameter model trained with the same tokenizer on Wikipedia and part of C4-en for one epoch.

C. Pre-training Details of the Llama3.1 Draft Model

D. Different Sampling Method

The section describes several sampling algorithms, reports TOKENSWIFT’s robustness across sampling methods, and explains tree-based attention’s candidate-branch design.

D.1. Introduction of Different Sampling Algorithms

  • Top-p sampling selects the smallest set whose cumulative probability exceeds p, whereas min-p retains tokens above pbase multiplied by the maximum token probability.Top-p samples from the reduced set using renormalized probabilities; min-p defines its pool using the scaled threshold.Source: §D.1 · p. 20
  • η-sampling defines its sampling pool as tokens whose probabilities meet or exceed η.The passage identifies entropy and α and ε as quantities in the η-sampling expression.Source: §D.1 · p. 20

E. Tree-Based Attention

  • Tree-based attention forms branches from combinations of candidate predictions and lets each token attend only to its predecessors.This mechanism processes multiple candidate continuations rather than selecting just one.Source: E. Tree-Based Attention · pp. 20–21

D.2. Impact of Different Sampling Algorithms

  • TOKENSWIFT demonstrates strong robustness across top-p, min-p, and η-sampling when token reutilization is disabled.The results are summarized in Table 12.Source: §D.2 · p. 20, Table 12 · p. 21

F.1. Ablation of Temperature

The temperature ablation evaluates top-p sampling at p = 0.9 across generation lengths from 20K to 100K tokens. Higher temperatures generally reduce acceptance, while speedup remains relatively stable or decreases slightly.

F.1. Ablation of Temperature

  • The ablation uses top-p sampling with p = 0.9 and tests temperatures from 0.4 to 1.2 over generation lengths of 20K to 100K tokens.Table 13 presents generation length, acceptance rate, and speedup results.Source: §F.1 · p. 21, Table 13 · p. 22
  • For 20K generation, acceptance rate falls from 0.79 at temperature 0.4 to 0.52 at temperature 1.2.A similar decrease is observed for longer sequences; speedup tends to remain relatively stable or decrease slightly as temperature rises.Source: §F.1 · p. 22
  • The highest speedups, around 3.4, occur across generation lengths at temperatures around 0.6 and 1.0.The passage characterizes moderate temperatures as offering the best balance between speed and quality.Source: §F.1 · p. 22, Table 13 · p. 22

F.2. Ablation of Prefill Length

The ablations examine how prefill length and penalty-window size affect speed, acceptance, and diversity. Longer prefill is associated with greater acceleration, while larger penalty windows trade some efficiency and acceptance for diversity.

F.2. Ablation of Prefill Length

  • A prefix length of 2048 yields 1.41 speedup for 20K tokens, while 8192 reaches up to 1.77 for 100K tokens.Acceptance rates generally remain around 0.35 to 0.39 across sequence lengths and prefix lengths.Source: Table 14 · p. 22, §F.2 · p. 22

F.3. Ablation of Penalty Window

  • Larger penalty windows, such as W = 256 or W = 2048, improve diversity while slightly reducing acceptance and maintaining significant speedup relative to AR.Smaller windows such as W = 20 yield higher acceptance and acceleration but lower Distinct-n diversity.Source: Table 15 · p. 23, Table 16 · p. 23, §F.3 · p. 23
  • For each penalty-window size, longer sequences from 20K to 100K generally bring higher acceleration and better diversity, with some acceptance-rate fluctuations.Source: Table 15 · p. 23, §F.3 · p. 23

H. Training Loss Curve

The supplied figure captions identify a training-loss curve for linear layers and a separate repetition-focused case study.

H. Training Loss Curve

  • Figure 6 presents generated-text fragments for YaRN-LLaMA2-7b-128k without and with Contextual Penalty.Source: Figure 6 · p. 24
  • In Figure 6, blue text marks the repeated portions of the generated text.Source: Figure 6 · p. 24
  • Figure 7 is a cross-entropy loss training curve for linear layers.Source: Figure 7 · p. 25
Loading 2502.18890v2…