Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

Speculative decoding needs a separate draft model, which is costly to acquire, train, and serve; maintaining a second model in memory also inflates memory footprint.

Add K lightweight decoding heads on top of the frozen backbone, each predicting token t+i directly from the backbone's last hidden state (no draft model, no separate draft weights). At each step the heads emit candidate continuations that are assembled into a tree; a tree-based attention mask lets the backbone verify all tree nodes in ONE forward pass, and a typical-acceptance scheme (bounded deviation from the target's own distribution, not strict argmax) decides which path survives. Medusa-1 fine-tunes heads on a frozen backbone (lossless); Medusa-2 fine-tunes heads jointly with the backbone for higher speedup at the cost of a special training recipe. A self-distillation variant covers the no-training-data case.

Medusa-1 over 2.2x, Medusa-2 2.3x-3.6x, on models of various sizes, versus plain autoregressive decoding. Model pairs are same-family draft heads (not a separate small model) with Vicuna/LLaMA-family backbones. The acceptance rate per step is the load-bearing number and is what should be reported alongside any speedup.

ADOPTION, precisely: PARTIAL engine support, not just research. vLLM ships Medusa support (medusa heads as a speculative method), SGLang ships Medusa-1.1, and llama.cpp has a self-speculative / Medusa-style draft path. But 'in-upstream-engine' is weaker than 'in-production' for the original recipe because Medusa-2's joint fine-tuning is not how anyone serves a foundation checkpoint - serving engineers use Medusa-1 heads trained separately for their own checkpoint. VENUE NOTE: no peer-reviewed venue verifiable via Semantic Scholar; arXiv comment carries a code URL only, so venue = arxiv-preprint and year = arXiv v1 (Jan 2024). Terminology trap: Medusa-1 and Medusa-2 are DIFFERENT recipes (frozen vs joint backbone), not version 1 and 2 of the method, and their speedups are not comparable. DANGLING REF RESOLVED 2026-10-04 by reading the paper (arXiv:2401.10774v1), and it CORRECTS THE OLD SHORTHAND'S FORM FACTOR. The paper names the PCIe card every time: "5 hours for MEDUSA-1 on Vicuna 7B model with a single NVIDIA A100 PCIE GPU" (section 2.2.1); "across various GPUs, including the A100-80GB-PCIe, A40, and A6000" (Appendix G.1); "Roofline Model (Llama 13B, A100 80GB PCIe)" (Figure 10); "We choose Llama 33B on A100-80GB-PCIe as the setting" (Appendix G.2); and the appendix tables are captioned "for Llama 33B on an A100 80GB PCIe". Source: https://arxiv.org/pdf/2401.10774 . So the 80 GB figure alone would have been ambiguous between the SXM4 and PCIe records, but the paper states PCIE explicitly: this is nvidia-a100-80gb-pcie, and the old 'nvidia-a100-sxm' value was wrong on form factor, not merely on suffix. Note the roofline figures quoted alongside (1,935GB/s HBM, 312 TFLOP/s) are the A100 80GB numbers in both cuts, so the memory-bandwidth figure alone would NOT have distinguished them. AFFILIATIONS BACKFILLED 2026-10-04 (paper-fields pass, agent w3p): OpenAlex institutions field, read per-author: Princeton (Cai, Lee, Dao), UIUC (Li, Chen), CMU (Geng), UConn (Peng). Filled from the API/PDF rather than from memory, because an affiliation is a fact about a specific paper's author block and inference from the author's employer history is how wrong affiliations get shipped. VENUE (+YEAR) CORRECTION AND VENUE_TRACK 2026-10-04 (paper-fields pass, agent w3p): venue 'arxiv-preprint' -> 'icml', venue_track 'main'. EVIDENCE: PMLR v235 (ICML 2024) table of contents contains the title verbatim: 'Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads'. v235 confirmed as the 41st ICML, 21-27 July 2024, from the volume landing-page header. BOTH VALUES RETAINED per the repo contract: the prior value was 'arxiv-preprint'. status set to contested because this changes a peer-reviewed attribution. Method note: OpenReview's api2 hard-429s under load and pre-2024 venues are unreachable through it, so these were verified against static proceedings indexes (proceedings.mlr.press for ICML, proceedings.mlsys.org for MLSys) rather than the OpenReview API. 2026-10-10: PMLR v235 ICML 2024 proceedings (idx2.py, pmlr.py: https://proceedings.mlr.press/v235/) confirms "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads". arXiv API 2401.10774 (axc.py) verified. Source niche-medusa-paper cites technical report vs PMLR v235 ICML 2024. Sources disagree; status remains contested.