Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

Medusa adds extra decoding heads on top of backbone LLM hidden states to predict multiple subsequent tokens in parallel. No separate draft model needed. Medusa-1 (frozen backbone): >2.2x speedup. Medusa-2 (joint fine-tuning): 2.3-2.8x. Uses tree-based attention for parallel verification. Typical acceptance scheme boosts acceptance rate at cost of generation quality. CRITICAL BATCH FINDING: 'when the batch size exceeds 32, the speedup decreases and may even have a negative effect' because linear layers shift from memory-bandwidth-bound to compute-bound. Experiments primarily at batch size 1. Accepted at ICML 2024.