Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads¶

Core mechanism on public data, not production kernel performance.

Paper information¶

Field	Details
Paper	arXiv 2401.10774
Institution	Together AI / Princeton University / University of Illinois Urbana-Champaign.
First public release	2024-01-19 (arXiv v1)
Author code	FasterDecoding/Medusa.
Local adapter or method key	`medusa`
Local reproduction code	`src/auto_research/reproductions/medusa/`

Paper method and evidence¶

Add future-token heads to a frozen or jointly fine-tuned backbone. Verify candidate branches together with tree attention to reduce serial decoding steps.

flowchart LR
  H[Backbone state] --> F[Multiple future-token heads]
  F --> T[Candidate tree]
  T --> V[Parallel backbone verification]

Original paper figure¶

The original figure shows multi-head proposals and tree verification. From the original paper; copyright remains with the authors.

Core equation¶

\[ \mathcal L=\sum_{k=1}^K\lambda_k\operatorname{CE}(p_k(x_{t+k}|h_t),x_{t+k}). \]

Reported results¶

The paper reports over 2.2× speedup for Medusa-1 and 2.3–3.6× for Medusa-2 while retaining generation quality.

Local reproduction¶

Compare single-head greedy with three future heads and tree verification: local backbone_calls 160.0000 → 80.0000 (-50.00%). Fixed metrics.

auto-research reproduce --paper medusa --dataset-dir data --seed 42

The WikiText-2 token transition model exercises future heads, candidate trees and backbone checks; no GPU tree-attention kernel or measured paper-scale latency.