Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

In plain words

Speculative decoding speeds up text generation by letting a fast helper guess ahead while the big model checks the guesses. The catch is that you need a separate, smaller helper model that speaks the same vocabulary, and running two models is awkward.

Medusa removes the helper. It bolts a few small extra prediction heads onto the big model itself. Where the normal output head predicts the next token, the extra heads predict the tokens two, three and more positions ahead. Their top guesses are combined into a small tree of candidate continuations, and the model checks the whole tree in a single pass, keeping the longest branch it agrees with.

The heads can be trained cheaply while the original model stays frozen, which leaves its outputs unchanged. Alternatively, they can be trained together with the model for more speed at a small cost in fidelity. The authors report generation two to almost three times faster for chat models at batch size one.

TL;DR

Medusa is draft-model-free speculative decoding. You attach K small "Medusa heads" (residual SiLU block plus vocab projection) to the backbone's final hidden state. Head k predicts the token k+1 positions ahead. Their top-k predictions form a fixed sparse tree of about 64 candidate tokens, and the backbone verifies the whole tree in one forward pass using a tree-shaped attention mask. Heads can be trained on a frozen (even quantized) backbone (Medusa-1, output unchanged under greedy or rejection-sampling acceptance) or jointly with a LoRA'd backbone (Medusa-2, higher speedup but the backbone changes). Acceptance is either standard rejection sampling or a new lossy "typical acceptance" rule. The authors report 2.18x to 2.83x wall-time speedup on Vicuna-7B/13B at batch size 1 in an HF-based setup. Speedup shrinks at larger batch sizes.

Problem

Autoregressive decoding is memory-bandwidth bound: each step moves all weights from HBM to produce one token, so compute sits idle §1. Speculative decoding helps, but it needs a separate draft model. That draft model is hard to obtain (separate pre-training, e.g. 275 A100 GPU hours reported for SpecInfer), can drift from the target distribution, and is awkward to serve in distributed systems §1 §2.1.1. Medusa assumes you can fine-tune small heads on data matching the model's output distribution, or generate such data with the model itself §2.2 §2.3.2. Experiments target batch size 1 (local, single-user serving) §1 §4.

Method

Each decoding step has three substeps, following speculative decoding: generate candidates (Medusa heads), process candidates (tree attention), accept candidates (rejection sampling or typical acceptance) §2. The heads sit on the backbone, so the logits computed while verifying step i also give the head predictions for step i+1. No separate drafting pass is needed §2.

1. Medusa heads. For last hidden state h_t, head k \in \{1..K\} predicts token t+k+1. The original LM head predicts t+1 and is called head 0 §2.1.1: p_t^{(k)} = \mathrm{softmax}\big(W_2^{(k)}(\mathrm{SiLU}(W_1^{(k)} h_t) + h_t)\big),\quad W_1^{(k)}\in\mathbb{R}^{d\times d},\ W_2^{(k)}\in\mathbb{R}^{d\times V} §2.1.1 Init: W_2^{(k)} is a copy of the original LM head and W_1^{(k)}=0, so each head initially reproduces the base model's next-token prediction §2.1.1.
2. Candidate tree. Take the top-s_k tokens from head k. The plain version is the Cartesian product across heads, a tree with s_k branches at level k (level 0 is the LM-head token). Total new tokens: \sum_{k=1}^K\prod_{i=1}^k s_i §2.1.2 fig. 2.
3. Tree attention. Flatten all tree nodes into one sequence and append them after the KV cache. Use a mask where each node attends only to the cached prefix and its own ancestors. Set each node's position id to prefix length + depth. One forward pass verifies all candidates without growing the batch §2.1.2 fig. 2.
4. Optimized sparse tree. On a calibration set, measure a_k^{(i)}, the accuracy of the i-th top prediction of head k alone (top-i minus top-(i-1) accuracy). A node for path [i_1..i_k] has estimated accuracy \prod_j a_j^{(i_j)} (assumes independence), and expected acceptance length is the sum over all nodes. Build the tree greedily: repeatedly add the connected node with the highest accuracy until the node budget is reached §2.3.3. The released tree has 64 nodes and depth 4, pruned using Alpaca-eval statistics app. C fig. 6.
5. Acceptance.
   • Rejection sampling as in speculative decoding gives the original model's distribution but does not raise the acceleration rate §1 §2.
   • Typical acceptance: always accept the first token, which is greedy from the LM head. Each later candidate token x_{n+k} is accepted while p_{\text{original}}(x_{n+k}\mid x_{\le n+k-1}) > \min\big(\epsilon,\ \delta\exp(-H(p_{\text{original}}(\cdot\mid x_{\le n+k-1})))\big) where \epsilon is a hard threshold and \delta an entropy-dependent one. The step's output is the longest accepted prefix over all candidates §2.3.1. At temperature 0 this reduces to greedy decoding §2.3.1.
6. Medusa-1 training (frozen backbone). Loss eq. 1: \mathcal{L}_{\text{Medusa-1}} = \sum_{k=1}^K -\lambda_k \log p_t^{(k)}(y_{t+k+1}) with \lambda_k = 0.8^k §2.2.1 app. B.2. The backbone only supplies hidden states, so it can be quantized §2.2.1.
7. Medusa-2 training (joint). Loss eq. 2: \mathcal{L}_{\text{Medusa-2}} = \mathcal{L}_{\text{LM}} + \lambda_0\,\mathcal{L}_{\text{Medusa-1}},\quad \mathcal{L}_{\text{LM}} = -\log p_t^{(0)}(y_{t+1}) It also uses a higher learning rate for the heads than the backbone and "heads warmup". Heads warmup means either two-stage training (Medusa-1 first, then joint) or gradually increasing \lambda_0 §2.2.2. The section's wording calls \lambda_0 "the weight of the backbone model's loss", but in eq. 2 it multiplies the heads' loss §2.2.2.
8. Self-distillation (no matching SFT data, or the model was RLHF'd). Take prompts from a seed set (ShareGPT, UltraChat) and let the model write the replies. Medusa-1 trains on this data directly. For Medusa-2, replace \mathcal{L}_{\text{LM}} with KL(p^{(0)}_{\text{original},t}\,\|\,p^{(0)}_t). The teacher is the same model with its LoRA adapter turned off, so no second model sits in memory. Use non-quantized LoRA here, otherwise the teacher is the quantized model §2.3.2.

Algorithm: Medusa decoding step (batch size 1)   [§2] [§2.1.2] [§2.3.1] [code: medusa/model/medusa_model.py#L314-L369]
Input: KV cache for prefix x_1..x_n, LM-head logits and K Medusa-head logits at the last accepted position, static tree T
1. t0 <- argmax (or sample) of the LM-head logits; for each head k take its top-TOPK tokens
2. Fill tree nodes: the node for path (i_1..i_k) gets the i_k-th top token of head k; the root gets t0
3. Forward all |T| nodes in one pass with the tree mask, position ids = n + depth
4. For each root-to-leaf candidate path, compare the backbone's prediction at each node with the next node's token
   (greedy: exact argmax match; typical: threshold rule above); accept_len = length of the cumulative-product prefix
5. Pick the candidate with the longest accept_len (typical/fast: tie-break by summed log-prob)
6. Append the root plus the accepted tokens; copy their KV entries into place; keep that path's LM/Medusa logits for the next step
Output: 1 + accept_len new tokens

Hyperparameters

Results

These numbers are author-reported. Setup: MT-Bench, batch size 1, baseline is the default HuggingFace implementation §3 §3.1. GPU for the speed measurements: not stated. Quality is the MT-Bench GPT-4 score (0-10) §3.2.

Definitions: speedup = acceleration rate / overhead, where overhead is per-step latency relative to vanilla decoding app. B.1.

Typical acceptance: raising \epsilon trades acceleration rate for quality. At high \epsilon, quality is comparable to random sampling §3.3.2 fig. 5.

Roofline/simulation (Llama-7B, A100): modeled acceleration is 0.477\log(\#\text{candidates}). Simulated speedup peaks at around 64 candidate tokens and falls once batch size exceeds 32 app. G.3 fig. 21 fig. 22.

Paper vs code

Checked against FasterDecoding/Medusa at e2a5d20.

• Head output projection. The paper gives each head its own W_2^{(k)}, initialized from the LM head §2.1.1. The current model gives each head a fresh nn.Linear(hidden, vocab, bias=False) with no copy from lm_head, so the weights must come from a checkpoint medusa_model.py. The legacy model used by the in-repo trainer has heads made only of ResBlocks and reuses the shared, frozen base lm_head for every head medusa_model_legacy.py medusa_model_legacy.py.
• Bias. The paper's head equation has no bias §2.1.1. ResBlock uses nn.Linear(hidden, hidden) with the default bias, and only the weight is zero-initialized medusa_model.py.
• Loss weights. The paper uses \lambda_k=0.8^k app. B.2. The in-repo trainer sums per-head cross-entropy without weights train_legacy.py. The label shift matches the paper: head i (0-indexed) is trained on labels shifted by 2+i train_legacy.py.
• Training recipe. The paper uses 5 heads, Vicuna v1.5, 8-bit AdamW, and 2 epochs app. B.2 §3.1. The repo scripts use Vicuna v1.3 with 3 heads, LR 1e-3, warmup_ratio 0.1, 1 epoch, and max length 2048 train_vicuna_7b.sh. The scripts call medusa/train/train.py, which does not exist at this commit (only train_legacy.py).
• Medusa-2 / self-distillation training is not in this repo. The README points to an Axolotl fork README.md.
• Rejection sampling is not implemented. Acceptance modes are greedy (temperature 0), typical, and nucleus utils.py.
• Typical acceptance details (default fast=True). The root token is argmax even at temperature > 0. Candidate probabilities come from temperature-scaled softmax. Entropy is computed with log(p + 1e-5). Ties among longest prefixes are broken by summed log-prob utils.py utils.py. The non-fast typical path calls get_typical_posterior_mask with 6 arguments, but the function takes 5, so it raises a TypeError utils.py utils.py.
• Default argument order. generate_candidates/evaluate_posterior default to posterior_threshold=0.3, posterior_alpha=0.09, the reverse of medusa_generate's 0.09/0.3 utils.py utils.py medusa_model.py.
• Batch size. Inference asserts batch size 1 medusa_model.py.

Pitfalls

• The headline range differs between versions. The arXiv metadata abstract says Medusa-2 reaches 2.3-3.6x, while the LaTeX abstract and §4 say 2.3-2.8x. The 3.62x figure is a single MT-Bench category (extraction) §3.1. Use about 2.8x as the overall number.
• Medusa-2 and typical acceptance are lossy. Medusa-2 changes backbone weights; tab. 1 shows MT-Bench deltas from -0.14 to +0.05. Typical acceptance does not preserve the sampling distribution. Only Medusa-1 with greedy or rejection-sampling acceptance is lossless [abstract] §2.3.1.
• Do not fine-tune the backbone directly with the heads. Quality dropped to 5.925 vs 6.17 tab. 2. Use two-stage training or a \lambda_0 warmup §2.2.2.
• Self-distillation for Medusa-2 needs KL to the original model's distribution, not hard labels. Training the backbone on self-generated text with CE alone degrades quality §2.3.2. With QLoRA, the "adapter-off" teacher is the quantized model §2.3.2.
• Tree mask and position ids must agree. Each node's position is prefix length + depth, not its index in the flattened sequence §2.1.2 utils.py. After acceptance, the accepted path's KV entries must be compacted into contiguous cache slots medusa_model.py.
• Bigger trees are not free. Acceleration grows roughly logarithmically with tree size while per-step cost grows. Past about 64 candidates, wall-clock speed falls §3.3.1 fig. 4 app. G.3. Optimal trees are model-specific (separate presets per model and stage in the code) and depend on calibration data app. C medusa_choices.py.
• The tree search assumes head accuracies are independent across positions §2.3.3. This is an approximation.
• Batch size 1 only in the experiments. Simulation predicts the speedup shrinks or turns negative above batch 32, as linear layers become compute-bound §4 app. G.3.
• Unstated: GPU and precision for the speed measurements, the sampling temperature for the main speedup numbers, the \epsilon,\delta used for the headline results, the warmup and epochs for self-distillation runs, and the head LR for self-distillation.

Lineage

• Builds on: Blockwise parallel decoding (Stern et al. 2018), which supplied the multi-head future-token prediction with residual FFN heads §2.1.1. Also speculative decoding and speculative sampling, whose draft/verify/accept framework and rejection sampling Medusa reuses with heads in place of the draft model §2. Tree verification parallels SpecInfer and staged speculative decoding §2.1.2 app. A. Typical acceptance adapts the truncation-sampling criterion of Hewitt et al. 2022 §2.3.1. Training uses LoRA and QLoRA app. B.2.
• Compared against: speculative decoding with open-source draft models (Llama-68M/160M, Tiny-Llama, Tiny-Vicuna) app. D tab. 1.
• Later work: EAGLE, a feature-level draft head.