Medusa

Overview

Medusa is a technique that accelerates LLM inference by predicting multiple future tokens in parallel using separate decoding heads, then verifying them. Achieves 2-3x speedup without model changes or distillation.

• Paper: "Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads" (Cai et al., 2023)
• Key Idea: Add auxiliary prediction heads for next N tokens
• Speedup: 2-3x with 1-2% quality loss
• Advantage: Works with any existing LLM without modification
• Trade-off: Increased inference parameters (auxiliary heads)

-

The Problem: Token-by-Token Generation is Sequential

Inference Bottleneck

Standard Autoregressive Generation:

Token 1: model(input) → output_1 → sample → token_1
Token 2: model(prev + token_1) → output_2 → sample → token_2
Token 3: model(prev + token_1 + token_2) → output_3 → sample → token_3
...

Timeline:
Step 1: ════ Forward pass (50ms)
Step 2: ════ Forward pass (50ms)
Step 3: ════ Forward pass (50ms)
...
Total for 100 tokens: 100 × 50ms = 5 seconds

Problem:
- Each token takes 50ms (compute)
- Must wait for previous token to continue
- Can't parallelize!
- Throughput: ~20 tokens/sec (very slow!)

Question:
Can we predict multiple tokens before committing to them?

Why Parallel Prediction is Feasible

Key observation:
Next token is mostly determined by recent context

Example:
"The capital of France is P" → Next token: "aris" (very predictable)
"Once upon a time" → Next token: hard to predict

But:
We can make EDUCATED GUESSES for next 2-3 tokens!

Correct sequence: "Paris, France"
Guess 1: "P" (likely)
Guess 2: "aris" (likely, follows P)
Guess 3: "," (likely, follows "Paris")

If guesses correct:
- Process in one forward pass instead of 3!
- 3x speedup!

If guess wrong:
- Reject and resample
- Still faster than sequential (amortized)

---

How Medusa Works

Architecture

Standard LLM:
- ┌─────────────────────────────────────┐
 - Input Embeddings │
 - ┤
 - Transformer Layers (32 layers) │
 - ┤
 - Final Layer Norm │
 - ┤
 - LM Head (output vocabulary) │
 - ┤
 - Output: vocab_size logits │
 - ┘
 ↓
 Sample token t+1

Medusa Modification:
- ┌─────────────────────────────────────┐
 - Input Embeddings │
 - ┤
 - Transformer Layers (32 layers) │
 - ┤
 - Final Layer Norm │
 - ┬───────────┤
 - LM Head │ (original)
 - ┼───────────┤
 - Medusa Head 1 │ (parallel head)
 - [small MLP] │
 - Output: vocab_size │ (predict t+2)
 - ┼───────────┤
 - Medusa Head 2 │ (parallel head)
 - [small MLP] │
 - Output: vocab_size │ (predict t+3)
 - ┼───────────┤
 - Medusa Head 3 │ (parallel head)
 - [small MLP] │
 - Output: vocab_size │ (predict t+4)
 - ┴───────────┘
 ↓ ↓ ↓ ↓
 Token t+1 Token t+2 Token t+3 Token t+4
 (main) (guess) (guess) (guess)

Inference Process

Step 1: Forward pass (once)
 - Input: prompt + previously accepted tokens
 - Output: logits from main head + all Medusa heads
 - Time: 50ms (same as before!)

Step 2: Generate candidate tokens
 - Main head: sample token t+1 (highest probability)
 - Medusa head 1: sample token t+2 (highest probability)
 - Medusa head 2: sample token t+3 (highest probability)
 - Medusa head 3: sample token t+4 (highest probability)
 - Time: <1ms (sampling is fast)

Step 3: Verification
 - Check: Is t+1 likely? (Yes, from main head)
 - Check: Given t+1, is t+2 likely?
 - Forward pass: model(prompt + t+1 + t+2_guess)
 - Compare main head output to t+2_guess
 - If probability > threshold: ACCEPT t+2
 - Else: REJECT, resample
 - Check: Given t+1, t+2, is t+3 likely?
 - Check: Given t+1, t+2, t+3, is t+4 likely?
 - Time: Usually ~50ms for 1-2 verifications

Result:
 - If all 4 tokens accepted: 4 tokens in 50ms (4x speedup!)
 - If 2 tokens accepted: 2 tokens in 50ms (2x speedup)
 - Average: 2-3x speedup

Real sequence vs guesses:
 - Model predicts: [high prob for "The"]
 - Medusa guesses: [high prob for "The"]
 - ACCEPT "The"
│
 - Model predicts (given "The"): [high prob for "cat"]
 - Medusa guesses: [high prob for "cat"]
 - ACCEPT "cat"
│
 - Model predicts (given "The cat"): [high prob for "sat"]
 - Medusa guesses: [high prob for "jumped"]
 - MISMATCH! REJECT jump, sample "sat"
│
Total: 3 tokens generated
- With guessing: 50ms + verification
- Without guessing: 150ms

Implementation

class MedusaHead(nn.Module):
 """Auxiliary decoding head for predicting future token"""

 def __init__(self, hidden_dim, vocab_size):
 super().__init__()
 self.head = nn.Sequential(
 nn.Linear(hidden_dim, hidden_dim),
 nn.ReLU(),
 nn.Linear(hidden_dim, vocab_size)
)

 def forward(self, x):
 return self.head(x)

class ModelWithMedusa(nn.Module):
 def __init__(self, base_model, hidden_dim, vocab_size, num_medusa_heads=3):
 super().__init__()
 self.base_model = base_model
 self.lm_head = base_model.lm_head

 # Medusa heads for parallel decoding
 self.medusa_heads = nn.ModuleList([
 MedusaHead(hidden_dim, vocab_size)
 for _ in range(num_medusa_heads)
])

 def forward(self, input_ids):
 """
 Returns:
 - main_logits: (batch, seq_len, vocab_size)
 - medusa_logits: list of (batch, seq_len, vocab_size)
 """
 # Base model forward
 hidden = self.base_model(input_ids, output_hidden_state=True)
 last_hidden = hidden.last_hidden_state

 # Main head
 main_logits = self.lm_head(last_hidden)

 # Medusa heads (parallel predictions)
 medusa_logits = [head(last_hidden) for head in self.medusa_heads]

 return main_logits, medusa_logits

def speculative_decoding_with_medusa(model, prompt, max_tokens=100, 
 medusa_temperature=0.5, 
 verify_threshold=0.9):
 """
 Generate with Medusa (speculative decoding)

 Key parameters:
 - medusa_temperature: Lower = more confident guesses
 - verify_threshold: Minimum probability to accept guess
 """
 generated = []
 current_input = prompt

 for _ in range(max_tokens):
 # Forward pass (single, not multiple!)
 main_logits, medusa_logits = model(current_input)

 # Latest token predictions
 main_pred = main_logits[:, -1,:] # (batch, vocab)
 medusa_preds = [m[:, -1,:] for m in medusa_logits]

 # Generate candidates
 main_token = sample(main_pred, temperature=0.7)
 medusa_tokens = [
 sample(m, temperature=medusa_temperature)
 for m in medusa_preds
]

 # Verify candidates
 candidates = [main_token] + medusa_tokens
 accepted_count = 0

 for i, token in enumerate(candidates):
 if i == 0:
 # Main token: always accept
 generated.append(token)
 accepted_count += 1
 else:
 # Verify: recompute with candidate token
 test_input = current_input + [generated[-1], token]
 test_logits, _ = model(test_input)

 # Get probability of this token
 token_prob = softmax(test_logits[:, -1,:])[0, token]

 if token_prob > verify_threshold:
 generated.append(token)
 accepted_count += 1
 else:
 # Reject and stop verifying
 break

 # Update input for next iteration
 current_input = current_input + generated[-accepted_count:]

 # Typically accept 2-3 tokens per forward pass
 # So 100 tokens takes ~30-50 forward passes instead of 100

 return generated

-

Performance Analysis

Speedup Calculation

Scenario: Generate 100 tokens with Medusa (3 heads)

Standard generation:
 - 100 forward passes × 50ms = 5000ms
 - Throughput: 20 tokens/sec

Medusa with 70% acceptance rate:
 - Main token: Always accepted
 - Head 1: 70% acceptance rate
 - Head 2: 50% acceptance rate (70% × 70%)
 - Head 3: 35% acceptance rate (70% × 70% × 70%)
 - Expected tokens per pass: 1 + 0.7 + 0.5 + 0.35 = 2.55 tokens/pass
 - Forward passes needed: 100 / 2.55 ≈ 39 passes
 - Verification passes: ~39 × 2 ≈ 78 passes (rough estimate)
 - Total passes: ~117 (vs 100 standard)
 - But each pass is parallel prediction, not sequential
 - Effective speedup: ~2.5x!
 - Throughput: 50 tokens/sec (2.5x improvement)

Better scenario (80% acceptance):
 - Expected tokens per pass: 1 + 0.8 + 0.64 + 0.51 = 2.95
 - Forward passes: 34
 - Speedup: ~3x
 - Throughput: 60 tokens/sec

Quality Impact

Experiment: LLaMA 7B with Medusa

Metric Baseline Medusa Delta
─────────────────────────────────────────────
MMLU Score 45.3% 44.8% -0.5%
Human Preference — Tied Equal
Generation Speed 1x 2.5x +150%

Quality loss:
 - Minimal (<1%)
 - Acceptance threshold tuning can improve
 - Trade: Fast inference worth slight quality dip

Configuration impact:
 - 3 heads: 2.5x speedup, 0.5% loss
 - 2 heads: 1.8x speedup, 0.2% loss
 - 1 head: 1.3x speedup, <0.1% loss
 - More heads = more speedup but potentially less accurate guesses

---

Advantages and Limitations

Advantages

2-3x speedup without model changes
Compatible with existing models
Minimal quality loss (<1%)
Works with batching
Orthogonal to other optimizations (can combine with quantization, etc.)

Limitations

Requires retraining Medusa heads
Quality degrades for "branching" outputs (multiple valid continuations)
Verification adds some compute overhead
Accepts only when confident (sensitive to threshold)
Less effective with high temperature (random sampling)

---

When to Use Medusa

Use Medusa when:
Inference speed is critical
Can afford to retrain auxiliary heads
Outputs are relatively deterministic
Want simple, model-agnostic solution

Avoid Medusa when:
Quality must be perfect
Outputs are highly variable (creative writing)
Can't modify inference code

---

Key Takeaways

Parallel prediction: Guess next tokens in parallel Verify guesses: Expensive but amortized over guesses 2-3x speedup with minimal quality loss Orthogonal to other optimizations Trade: Slight quality for major speed gain

-

Related Notes

• Speculative Decoding - Similar concept, different approach
• Continuous Batching - Works together for inference optimization
• Token Merging (Tome) - Another inference acceleration technique
• Llm Inference Optimization - Complete inference stack