PaperlayerView exact revision on arXiv

Source-linked AI summary

Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao

arXiv:2401.10774v3cs.LGcs.CL

TL;DR

LLM inference is slowed by memory-bandwidth-bound, sequential autoregressive decoding, while speculative decoding requires a difficult-to-obtain separate draft model. Medusa adds decoding heads that generate and tree-verify multiple candidate tokens in parallel, with two fine-tuning procedures and extensions for acceptance and data generation. Across models and prompt types, it achieves 2.3–2.8x speedup without compromising generation quality, while its evaluation focuses primarily on batch size one.

  • Problem

    LLM inference is memory-bandwidth-bound and sequential, while speculative decoding depends on a challenging separate draft model.

  • Method

    Medusa adds predictive decoding heads, tree-based parallel candidate verification, two fine-tuning procedures, typical acceptance, and self-distillation when training data is unavailable.

  • Results

    2.3–2.8x speedup is achieved across different prompt types and models without compromising generation quality.

  • Takeaways & Limitations

    Medusa provides a parameter-efficient, easily integrated alternative to draft-model-based speculative decoding for accelerating existing LLM systems.

  • Takeaways & Limitations

    The experiments primarily focus on batch size one, although the authors state that the ideas can generalize to larger batches.

Abstract

from arXiv · show

Large Language Models (LLMs) employ auto-regressive decoding that requires sequential computation, with each step reliant on the previous one's output. This creates a bottleneck as each step necessitates moving the full model parameters from High-Bandwidth Memory (HBM) to the accelerator's cache. While methods such as speculative decoding have been suggested to address this issue, their implementation is impeded by the challenges associated with acquiring and maintaining a separate draft model. In this paper, we present Medusa, an efficient method that augments LLM inference by adding extra decoding heads to predict multiple subsequent tokens in parallel. Using a tree-based attention mechanism, Medusa constructs multiple candidate continuations and verifies them simultaneously in each decoding step. By leveraging parallel processing, Medusa substantially reduces the number of decoding steps required. We present two levels of fine-tuning procedures for Medusa to meet the needs of different use cases: Medusa-1: Medusa is directly fine-tuned on top of a frozen backbone LLM, enabling lossless inference acceleration. Medusa-2: Medusa is fine-tuned together with the backbone LLM, enabling better prediction accuracy of Medusa heads and higher speedup but needing a special training recipe that preserves the backbone model's capabilities. Moreover, we propose several extensions that improve or expand the utility of Medusa, including a self-distillation to handle situations where no training data is available and a typical acceptance scheme to boost the acceptance rate while maintaining generation quality. We evaluate Medusa on models of various sizes and training procedures. Our experiments demonstrate that Medusa-1 can achieve over 2.2x speedup without compromising generation quality, while Medusa-2 further improves the speedup to 2.3-3.6x.

1. Introduction

MEDUSA addresses the memory-bandwidth and sequential-decoding bottleneck in LLM inference by adding parallel decoding heads without requiring a separate draft model. It offers two fine-tuning procedures and achieves 2.3–2.8x speedup across models and prompt types without compromising generation quality.

  • Motivation: LLM inference latency is dominated by memory bandwidth and the sequential nature of autoregressive decoding.Increasing model size improves generation quality but makes practical inference slower.
  • Motivation: Speculative decoding reduces decoding steps with a smaller draft model, but obtaining and maintaining an appropriate draft model is challenging.Separate draft-model training can require substantial computational resources and may introduce distribution shift.
  • MEDUSA: MEDUSA adds parameter-efficient decoding heads to an existing backbone model to predict multiple subsequent tokens concurrently, eliminating the need for a separate draft model.The heads can be integrated into existing systems, including distributed environments.
  • MEDUSA: MEDUSA generates multiple candidate continuations, verifies them concurrently through adjusted attention, and uses rejection sampling or typical acceptance to select continuations.The typical acceptance scheme can accelerate decoding further while maintaining similar generation quality.
  • Fine-tuning: MEDUSA-1 freezes the backbone for quality-preserving, resource-efficient fine-tuning, whereas MEDUSA-2 jointly fine-tunes the heads and backbone for greater speedup.MEDUSA-2 uses a training protocol intended to preserve next-token prediction capability and output quality; self-distillation can create training data when datasets are unavailable.
  • Results: 2.3–2.8x speedup is achieved across different prompt types and models without compromising generation quality.Experiments cover Vicuna-7B, Vicuna-13B, Vicuna-33B, and Zephyr-7B under varied training settings.

2. Methodology

MEDUSA accelerates autoregressive LLM inference by adding decoding heads, generating multiple candidate continuations, and verifying them concurrently with tree attention. It offers distinct training strategies and acceptance extensions for different data, resource, and quality requirements.

  • MEDUSA pipeline: MEDUSA heads generate subsequent-token predictions, while tree attention processes candidate continuations concurrently and candidate acceptance completes each decoding step.The pipeline generates candidates with MEDUSA heads, processes them with tree attention, and accepts them using rejection sampling or typical acceptance.
  • Tree attention: Tree attention uses predecessor-only masking and adjusted positional indices to process many candidates simultaneously without expanding the batch size.Candidates are formed from top-s_k predictions at successive levels; the Cartesian-product example yields 2 × 3 = 6 candidates.
  • Fine-tuning: MEDUSA-1 freezes the backbone and fine-tunes only the additional heads, preserving generation quality while requiring minimal memory for existing-model integration.This strategy is intended for limited computational resources or deployment without affecting the backbone model's performance.
  • Fine-tuning: Jointly training the backbone and MEDUSA heads can improve head accuracy, using training strategies designed to preserve the backbone model's capability.The recipe combines differential learning rates and heads warmup, and can be applied with supervised fine-tuning.
  • Extensions: Self-distillation supplies training data when the original dataset is unavailable, while typical acceptance selects plausible candidates using the original model's prediction probabilities.Typical acceptance evaluates candidate prefixes with hard and entropy-dependent thresholds, then chooses the longest accepted prefix.

3. Experiments

Experiments evaluate MEDUSA across model sizes, training settings, tree configurations, sampling schemes, and fine-tuning strategies. MEDUSA consistently accelerates inference while preserving generation quality, with speed–overhead and speed–quality trade-offs depending on configuration.

  • MEDUSA-1 v.s. MEDUSA-2 on Vicuna 7B and 13B: 2.18× and 2.33× speedups are achieved by MEDUSA-1 on Vicuna-7B and Vicuna-13B, respectively, while MEDUSA-2 reaches 2.83× on both models.The coding and extraction MT-Bench categories reach 3.29× and 3.62× speedups with MEDUSA-2 Vicuna-7B.
  • Configuration of Tree Attention: 64-node optimized sparse trees achieve better acceleration than 256-node dense trees, while increasing tree complexity eventually reduces speed because of computational overhead.Acceleration increases logarithmically and slows as tree size grows.
  • Thresholds of Typical Acceptance: As the typical-acceptance threshold increases, quality rises while acceleration decreases; at higher thresholds, typical sampling approaches random sampling for creative tasks.The comparison uses writing and roleplay on MT-Bench with MEDUSA-2 Vicuna-7B.
  • Effectiveness of Two-Stage Fine-Tuning: 2.18x speedup is obtained by fine-tuning only MEDUSA heads, whereas the two-stage MEDUSA-2 procedure reaches 2.83x while maintaining generation quality.Direct fine-tuning with MEDUSA heads degrades generation quality.

4. Discussion

MEDUSA accelerates LLM inference by adding predictive decoding heads that generate multiple tokens simultaneously, with typical acceptance simplifying sampling. The paper focuses experimentally on batch size 1, while noting broader applicability.

  • 2.3-2.8 times faster inference is achieved by equipping models with additional predictive decoding heads that generate multiple tokens simultaneously.The approach is described as simple, parameter-efficient, and easy to integrate into existing systems.
  • Typical acceptance removes complications from rejection sampling while providing reasonable outputs.
  • Batch size 1 is the primary experimental setting, although the authors state that the ideas can generalize to larger batch sizes.The paper identifies batch size 1 as representative of locally hosted personal-use scenarios.

Impact Statement

MEDUSA targets the memory-bandwidth-bound latency of autoregressive LLM inference by reducing decoding steps and improving arithmetic intensity. Its broader implications include accessibility and deployment benefits, alongside inherited bias, interpretability, misuse, and sampling-quality considerations.

  • Broader implications: MEDUSA’s efficiency may lower deployment barriers and broaden access to advanced AI capabilities across organizations and sectors.
  • Ethical considerations: MEDUSA inherits backbone-model concerns about bias and fairness, while its tree attention and multiple heads may complicate interpretability and accountability.
  • Ethical considerations: Acceleration may enable misuse such as scaled disinformation generation or automated cyber-attacks, motivating ethical guidelines and security measures.
  • Motivation: MEDUSA addresses memory-bandwidth-bound autoregressive inference, where model parameters and the KV cache consume substantial memory resources.
  • Related methods: Speculative decoding reduces decoding steps by having a smaller draft model propose several tokens for collective evaluation by the LLM.
  • Sampling: MEDUSA’s typical acceptance excludes improbable candidates without requiring exact correspondence to the language-model distribution, supporting diverse, high-quality outputs.

B.4. Training with Self-Distillation on Vicuna-33B and Zephyr-7B

The reported training setup uses MEDUSA-2 with a sine schedule for one model configuration and a two-stage procedure for another, while the sparse tree uses four MEDUSA heads and pruning based on top-k prediction statistics.

  • Training with self-distillation: MEDUSA-2 training can use a sine schedule for θ0, gradually increasing it to its peak by the end of training.The authors report this approach as equally effective for the specified models.
  • Sparse tree construction: The sparse Vicuna-7B tree extends four levels deep, corresponding to four MEDUSA heads used in computation.The tree is formed by Cartesian products and refined through pruning based on top-k prediction statistics from Alpaca-eval.

D. Results of Speculative Decoding

Speculative decoding performance depends on matching the draft model and draft-token count to the target Vicuna model size. The reported comparisons show different optimal configurations for 7B, 13B, and 33B models.

  • Evaluation setup: The evaluation applies speculative decoding to Vicuna models with 7B, 13B, and 33B parameters using open-source draft models.The draft models include Llama-68M, Llama-160M, Tiny-Llama, and Tiny-Vicuna.
  • Tree visualization: Figure 6 depicts a sparse MEDUSA-2 Vicuna-7B tree with 64 candidate-token nodes and depth 4, with red lines marking the correctly predicted future-token path.
  • Results: γ = 4 with Llama-68M performs best for Vicuna-7B, while γ = 3 with the same draft model is most effective for Vicuna-13B.
  • Results: γ = 3 with Tiny-Vicuna provides the greatest acceleration for Vicuna-33B.
  • Discussion: Draft-model settings must be tailored to the target LLM size, leaving configuration choice as an area for further exploration.

E. Additional Results for All Models

The paper reports speedup results across various models, summarized in Figure 8.

  • Figure 8 reports speedup results for various models with MEDUSA-2.

F. Additional Results on AlpacalEval Dataset

MEDUSA-2 maintains consistent speedup on AlpacaEval, while self-distilled models show weaker speedup when preserving quality trades off against speed.

  • MEDUSA-2 achieves consistent speedup on AlpacaEval, similar to its MT-Bench results.
  • MEDUSA-2 shows significant speed improvement across the evaluated models.Self-distilled Zephyr-7B and Vicuna-13/33B have weaker speedup because preserving quality trades off against boosting speed.
  • The hardware analysis profiles MEDUSA-style decoding across GPUs, model sizes, sequence lengths, and batch sizes.It combines operator benchmarks with an analytical model of acceleration rates.

G.1. Roofline Model of Operators

The roofline analysis benchmarks Llama operators across GPUs and shows that decoding attention and linear layers are constrained primarily by memory bandwidth.

  • G.1. Roofline Model of Operators: The study benchmarks Llama-7B, Llama-13B, and Llama-33B operators on A100-80GB-PCIe, A40, and A6000 GPUs.It focuses on three categories of matrix-multiplication operators as major computational-overhead sources.
  • G.1. Roofline Model of Operators: The analysis compares computational and space complexity for linear and attention operators across prefill, decoding, and MEDUSA decoding phases.The operators include query/key/value projections, attention matrix multiplications, and up/gate/down layers.
  • G.1. Roofline Model of Operators: 49 settings per operator vary batch sizes from 1 to 64 and sequence lengths from 128 to 8192.Figures 9–17 report these benchmarks across model sizes and GPUs.
  • G.1. Roofline Model of Operators: During decoding, attention matrix multiplications remain memory-bandwidth-limited as Operational Intensity stays nearly unchanged despite higher FLOP/s.The pattern indicates suboptimal resource utilization in self-attention.
  • G.1. Roofline Model of Operators: During decoding, linear-layer datapoints follow the GPU memory-bandwidth slope, indicating that these layers are also memory-bandwidth-bound.Increasing batch size improves achieved FLOP/s and Operational Intensity under this constraint.

G.2. FLOP/s vs. Operational Intensity Variations in MEDUSA

The MEDUSA hardware analysis examines how adding candidate tokens changes Operational Intensity and FLOP/s for attention and linear layers.

  • G.2. FLOP/s vs. Operational Intensity Variations in MEDUSA: The analysis evaluates Llama-33B on an A100-80GB-PCIe to study MEDUSA’s effects on operator efficiency.
  • G.2. FLOP/s vs. Operational Intensity Variations in MEDUSA: 44× FLOP/s and 41× Operational Intensity are achieved for attention matrix multiplication as MEDUSA adds candidate tokens at batch size 16.The comparison is against regular decoding.
  • G.2. FLOP/s vs. Operational Intensity Variations in MEDUSA: For up/gate/down linear layers, increasing candidate tokens and batch size shifts execution from a memory-bandwidth-bound region toward a computation-bound region.The linear-layer results are reported in Figure 20 and Table 8.
  • G.2. FLOP/s vs. Operational Intensity Variations in MEDUSA: For attention matrix multiplication with sequence length 1024, the analysis compares FLOP/s and Operational Intensity as MEDUSA varies candidate tokens.

G.3. Predicting MEDUSA Performance

MEDUSA performance depends on selecting an appropriate number of candidate tokens and balancing batch size and sequence length. Simulations show diminishing returns from excessive candidates, reduced speedup at large batches, and lower performance for longer sequences.

  • acc rate = 0.477 log(num candidate) approximates MEDUSA's acceleration rate in the analytical model.The model simulates a simplified Llama-7B block with batch size 1 and sequence length 1024, processing candidate tokens in parallel through tree attention.
  • Speedup initially improves as the number of candidate tokens increases but declines beyond 64 candidates.The decline indicates diminishing returns from further increases in candidate length and matches the reported experimental trend.
  • Speedup decreases and may become negative when batch size exceeds 32.At large batch sizes, linear layers shift from memory-bandwidth-bound to computationally bound.
  • The optimal number of candidate tokens remains relatively consistent across sequence lengths, while overall performance decreases as sequence length grows.The performance drop is primarily attributed to attention matrix multiplication overhead; linear-layer computation remains independent of sequence length.
  • MEDUSA scaling benefits require optimizing candidate count, batch size, and attention overhead together.Increasing batch size initially helps through parallelism, but excessive batching and longer sequences reduce speedup.
Loading 2401.10774v3…