PaperlayerView exact revision on arXiv

Source-linked AI summary

Break the Sequential Dependency of LLM Inference Using Lookahead Decoding

Yichao Fu, Peter Bailis, Ion Stoica, Hao Zhang

arXiv:2402.02057v1cs.LGcs.CL

TL;DR

Autoregressive LLM decoding is memory-bandwidth bounded and underuses accelerator parallelism, while draft-model methods are difficult to train and generalize. The paper introduces exact Lookahead decoding, which generates and verifies n-grams in parallel without auxiliary models, achieving up to 1.8x speedup on MT-Bench and 4x on 8 GPUs for code completion. Its central trade-off is fewer decoding steps through increased per-step computation, with performance depending on the hardware and serving regime.

  • Problem

    Autoregressive decoding has high latency from sequential token generation and underused accelerator parallelism, while draft models are non-trivial to obtain and do not generalize across base models and datasets.

  • Method

    Lookahead decoding is a lossless parallel algorithm that generates and verifies n-grams using lookahead and verification branches plus an n-gram pool, without auxiliary components.

  • Results

    The implementation achieves 1.8x speedup on MT-Bench and up to 4x speedup on 8 GPUs for code completion with Lookahead Parallelism.

  • Takeaways & Limitations

    Lookahead decoding can linearly reduce decoding steps by investing exponentially in per-step FLOPs while preserving the output distribution.

  • Takeaways & Limitations

    Extra computations can cause slowdowns in compute-bound environments, and lower speedup is observed on GPUs with smaller FLOPs capacity.

Abstract

from arXiv · show

Autoregressive decoding of large language models (LLMs) is memory bandwidth bounded, resulting in high latency and significant wastes of the parallel processing power of modern accelerators. Existing methods for accelerating LLM decoding often require a draft model (e.g., speculative decoding), which is nontrivial to obtain and unable to generalize. In this paper, we introduce Lookahead decoding, an exact, parallel decoding algorithm that accelerates LLM decoding without needing auxiliary models or data stores. It allows trading per-step log(FLOPs) to reduce the number of total decoding steps, is more parallelizable on single or multiple modern accelerators, and is compatible with concurrent memory-efficient attention (e.g., FlashAttention). Our implementation of Lookahead decoding can speed up autoregressive decoding by up to 1.8x on MT-bench and 4x with strong scaling on multiple GPUs in code completion tasks. Our code is avialable at https://github.com/hao-ai-lab/LookaheadDecoding

1. Introduction

Autoregressive LLM decoding is latency-limited because it generates one token per step while underusing accelerator parallelism. Lookahead decoding addresses this with exact, parallel n-gram generation and verification without an auxiliary draft model.

  • Autoregressive decoding generates one token per step and largely underutilizes modern accelerators, making low-latency long-sequence generation difficult.
  • Speculative decoding uses a draft model to guess tokens for parallel verification, but training an effective draft model is non-trivial and its generalization is limited.
  • Lookahead decoding uses lookahead and verification branches plus an n-gram pool to generate, verify, cache, and integrate subsequent-token candidates.
  • 1.8x speedup is achieved on MT-Bench, while up to 4x speedup is achieved for code completion with Lookahead Parallelism on 8 GPUs.
  • Lookahead decoding is a lossless parallel algorithm that accelerates LLM inference without requiring an auxiliary component.
  • Lookahead decoding linearly reduces decoding steps according to per-step log(FLOPs), enabling tradeoffs between per-step computation and total steps.
  • The implementation supports memory-efficient attention and distributed CUDA execution, and evaluations demonstrate effectiveness under different settings.

2. Background

The background formulates autoregressive decoding as sequential token prediction and Jacobi decoding as nonlinear-system iteration. Lookahead decoding builds on the resulting parallel trajectories while verification preserves accepted-token correctness.

  • Causal decoder attention ensures each output token depends only on the current and preceding input tokens.
  • Autoregressive decoding predicts each next token from all previous tokens, using sampling methods such as greedy, top-K, and top-P.
  • Under greedy sampling, generating an m-token output solves m token-selection problems sequentially from the prompt.
  • Speculative decoding exemplifies guess-and-verify by generating a draft sequence with a draft model and verifying its tokens in parallel with the LLM.
  • Lookahead decoding maintains a fixed two-dimensional trajectory window, caches generated n-grams, and verifies promising candidates before integrating them.
  • Jacobi decoding iteratively updates all token variables from an initial guess until reaching a fixed point, guaranteeing the full solution in at most m iterations.
  • Adjacent tokens from successive Jacobi trajectory states can form meaningful 2-grams, creating material for parallel n-gram generation.
  • Jacobi decoding alone often fails to reduce decoding steps because tokens appear in wrong positions or are replaced by later iterations.

3. LOOKAHEAD DECODING

LOOKAHEAD DECODING combines parallel n-gram generation with later verification, preserving the output distribution while reducing wasted compute. Its attention design and token-level workload partitioning support efficient execution on single and multiple GPUs.

  • 3.1. Lookahead Branch: LOOKAHEAD DECODING uses a fixed-sized 2D window to generate multiple disjoint n-grams from the Jacobi decoding trajectory.W controls lookahead into future positions, while N controls lookback into past trajectory steps.
  • 3.2. Verification Branch: The lookahead branch caches generated n-grams in a pool, while the verification branch progressively checks promising candidates before integrating accepted n-grams.Verification supports advanced sampling by removing n-grams with mismatched prefixes.
  • 3.2. Verification Branch: At most G candidates run in parallel in the verification branch, and practice sets G = W to balance n-gram generation with verification cost.The cap manages the growing verification workload as the n-gram cache expands.
  • 3.2. Verification Branch: A designated attention mask integrates lookahead and verification into one decoding step while preventing tokens in either branch from attending to the other.The mask is derived so each token is visible only to tokens with larger position indices than itself.
  • 3.4. Lookahead Parallelism: Lookahead parallelism distributes disjoint lookahead branches and independent verification candidates across GPUs without communication during the forward pass.Each GPU maintains an entire model copy, requiring more memory but enabling near-zero communication per decoding step.

4. Scaling Law of LOOKAHEAD DECODING

This section relates Lookahead decoding’s flexible per-step computation to its step compression ratio and compares its scaling behavior with speculative decoding. The formulation and experiments indicate that increasing parallel computation can reduce decoding steps, including when scaling across GPUs.

  • 4.2. Estimating Speedup: Speculative decoding models accepted tokens from b guessed sequences of length γ using token acceptance rate β with expectation E(β) = α.When b = 1, the multi-sequence formulation reduces to the single-sequence case.
  • 4.1. Step Compression Ratio: The step compression ratio S is the number of autoregressive steps divided by the Lookahead decoding steps needed to generate the same sequence length.It measures reduction in total decoding steps rather than a change in generated sequence length.
  • 4.2. Estimating Speedup: Lookahead decoding assumes one good speculation every f steps, accepting E(#tokens) tokens on that step and falling back to autoregressive decoding on the other f −1 steps.This assumption connects expected accepted tokens per step to the step compression ratio.
  • 4.2. Estimating Speedup: The empirical MT-Bench trend aligns to some extent with the formulation, which predicts a linear reduction in decoding steps with per-step log(b) when γ is sufficiently large.The paper contrasts this with speculative decoding’s upper bound under its formulation.
  • 4.2. Estimating Speedup: Lookahead decoding’s scaling law relates linear decoding-step reduction to per-step log(FLOPs) for sufficiently large N and suggests stronger latency reduction across multiple GPUs.Per-step FLOPs are roughly proportional to (W + G) ∗(N −1).

5. Evaluation Results

LOOKAHEAD DECODING is evaluated across models, datasets, attention implementations, distributed settings, and sampling conditions. It delivers substantial speedups, while requiring extra computation and surplus FLOPs that limit gains in compute-bound or lower-FLOP settings.

  • Experimental Setup: The evaluation spans LLaMA-2 and CodeLlama models from 7B to 70B across chat, math, code, summarization, and instruction-based generation tasks.Experiments use A100-based single- and multi-GPU setups, including an 8-GPU DGX system.
  • End-to-end Performance: 1.5x-2.3x speedups are achieved across datasets against HuggingFace greedy search, with code completion reaching 2.3x.Higher repetition in code completion makes predictions easier; smaller models generally show higher speedups because they have more available FLOPs.
  • LP and FlashAttention: 1.8x speedup is achieved for 7B MT-Bench with FlashAttention, while strong scaling with Lookahead Parallelism reaches 4x on ClassEval.FlashAttention speeds up the PyTorch implementation by 20%; TP and PP instead introduce 0.75x-0.82x slowdowns in the reported distributed setting.
  • Sampling: 1.46x-1.60x speedups are obtained on CNN/Daily Mail and XSum while preserving output-distribution quality under sampling and greedy search.Sampling produces smaller speedups because its acceptance ratio is lower.
  • Ablation Study: A balanced lookahead and verification design outperforms a configuration with a tiny verification branch, which has lower speculation acceptance potential.Prompt lookup can further boost LOOKAHEAD DECODING and is integrated into the implementation.
  • Discussion and Limitation: Extra computation is the main limitation: 7B, 13B, and 34B models require 120x, 80x, and 56x extra per-step FLOPs, respectively.High speedups require surplus FLOPs; compute-bound serving or GPUs with lower FLOP capacity can reduce or limit gains, although MT-Bench still shows 30% on RTX 3090 and > 50% on A100.

6. Related Work

Speculative decoding accelerates autoregressive decoding through a guess-and-verify process, using a smaller draft model to propose tokens that the base LLM verifies in parallel.

  • Speculative decoding uses a draft model to speculate several subsequent tokens before parallel verification by the base LLM.The draft model requires fewer resources, while verifying multiple tokens in parallel costs about as much as generating one token.

7. Conclusion

Lookahead decoding parallelizes autoregressive LLM decoding without changing the output distribution. It achieves speedups without a draft model and can reduce decoding steps through increased per-step FLOPs.

  • Lookahead decoding parallelizes autoregressive decoding without changing the output distribution.
  • Lookahead decoding can linearly decrease decoding steps with exponential investment in per-step FLOPs.

A. Algorithms

The algorithms combine Jacobi-style lookahead generation, cached n-gram speculation, and greedy or sampling-based verification within iterative decoding.

  • Jacobi decoding: Jacobi decoding initializes token positions and iteratively generates output tokens until a stopping condition is reached.
  • Lookahead decoding: Lookahead decoding accepts a prompt, model, n-gram size, window size, speculation limit, and maximum decoding steps.
  • Lookahead decoding: The lookahead branch generates windowed token sequences, while the verification branch selects n-grams from the pool and appends verified output.
  • Greedy verification: Greedy verification compares candidate speculation tokens with model argmax outputs, accepting matches and updating remaining potential speculations.
  • Sample verification: Sample verification samples acceptance decisions using candidate probabilities and updates the speculation pool after accepted tokens.

B. Proof: Output distribution preserved disjoint n-gram verification

The sampling-verification proof establishes that Lookahead decoding preserves the LLM’s output distribution under greedy-generated speculations by induction over the number of speculations.

  • Lookahead sampling verification uses greedy-generated speculations while preserving the LLM’s output distribution.
  • The proof defines P(v) as the LLM probability and Q(v) as the algorithm’s probability, then establishes P(v) = Q(v).
  • Induction proves Q_G(v) = P(v) for any number of speculations, beginning with the single-speculation case.
  • The rejection cases update probabilities so that Q_G(v) remains equal to P(v), including when the candidate differs from v or has zero probability.
  • The complete argument guarantees correctness of the sampling algorithm.

C. Derivation of Expectation of The Number of Accepted Tokens

The derivation computes the expected number of accepted tokens for single-candidate speculation and then extends the acceptance probabilities to a batch of b speculations.

  • Single-candidate speculation: Single-candidate speculation assigns zero probability to accepting i ≥ γ + 2 tokens.The maximum possible accepted-token count is γ + 1.
  • Single-candidate speculation: For i ≤ γ, accepting i tokens has probability equal to the probability of accepting i − 1 tokens multiplied by α.
  • Single-candidate speculation: The single-candidate expectation expands as 1 ∗(1 −α) + 2 ∗(1 −α) ∗α + ... + (γ + 1) ∗αγ.
  • Single-candidate speculation: The expectation is algebraically rearranged by grouping successive terms involving powers of α.
  • Batch speculation: For batch size b, pi denotes the probability that at most i tokens are accepted across all b speculations, with P(#accepted tokens = i) = pi −pi−1 for i ≤ γ.
  • Batch speculation: The batch expectation weights these probabilities through (p1 −p0) + (2p2 −2p1) + ... + (γ + 1)(1 −pγ).

D. Prompt for LLaMA-2-Chat on Summarization Tasks

The summarization task uses a prompt adapted from Ruan et al. (2023) that instructs LLaMA-2-Chat to answer only from the original text and produce direct, readable output.

  • The summarization prompt is modified from Ruan et al. (2023).
  • The system instructions require answers to use only the supplied original-text context.
  • The prompt requests human-readable output without gibberish and restricts generation to the requested output.
  • The instructions prohibit prefatory phrases such as thanking the user or identifying as an AI agent.

E. Verification of Generation Quality for Greedy Sampling and Advanced Supports

The verification evaluates whether FlashAttention and LP support preserve Lookahead Decoding’s generation quality, using compression ratio comparisons across models and tasks.

  • FlashAttention and LP support do not change the compression ratio S of vanilla Lookahead Decoding.
  • 18 generations across 7B and 13B models on MT-Bench, HumanEval, and ClassEval were compared with and without FlashAttention.
  • The average compression ratio S was 3.267 with FlashAttention and 3.259 without it, a difference below 0.3%.
Loading 2402.02057v1…