Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

By Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, De-huai Chen, Tri Dao. International Conference on Machine Learning, 2024

809 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, General Social Sciences

View Paper: https://www.semanticscholar.org/paper/57e7af0b69325fafb371ef5d502e39ef9c90ef7e

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In General Social Sciences

Citations in the corpus

BLOCK DIFFUSION: INTERPOLATING BETWEEN AU-TOREGRESSIVE AND DIFFUSION LANGUAGE MODELS (2025, neutral): “…that compete with diffusion include iterative editing (Gu et al., 2019), parallel and speculative decoding (Gu et al., 2017; Santilli et al., 2023; Cai et al., 2024; Gloeckle et al., 2024), consistency training (Kou et al., 2024), guidance (Sanchez et al., 2023), and cross-modal extensions (Liu…”
SnapKV: LLM Knows What You are Looking for Before Generation (2024, methodology): “Our investigation incorporates SnapKV with Medusa [35] 3 , a cutting-edge parallel decoding framework that utilizes multiple classifiers and tree attention mechanisms for drafting tokens, subsequently verified by LLMs.”
EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees (2024, methodology): “EAGLE (Li et al., 2024b) and Medusa (Cai et al., 2024) use the same static draft tree structure in all contexts: at the i -th step of the draft phase, k candidates are added, with k being fixed.”
Better & Faster Large Language Models via Multi-token Prediction (2024, methodology): “Stern et al. (2018) and Cai et al. (2024) propose model finetunings with multi-token prediction for faster inference but do not study the effects of such a loss during pretraining.”
QuIP#: Even Better LLM Quantization with Hadamard Incoherence and Lattice Codebooks (2024, background): “This not only lets large models fit on smaller devices, but also enables faster throughput in memory bound settings such as autoregressive decoding.”
Break the Sequential Dependency of LLM Inference Using Lookahead Decoding (2024, methodology): “Several approaches have been proposed – one such approach is speculative decoding (Chen et al., 2023; Leviathan et al., 2023) and its variants (He et al., 2023; Stern et al., 2018; Cai et al., 2024; Li et al., 2023; Liu et al., 2023; Miao et al., 2023).”