Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
Tianle Cai, Yuhong Li, Zhengyang Geng, et al.2024ICML 2024
Read it on arxiv.org(opens in a new tab)Why this one
Read this after Fast Inference from Transformers via Speculative Decoding. Leviathan gives you the clean idea: let a cheap model guess, let the big model verify. Medusa asks a very practical serving question: what if the draft model is the part you do not want to own? Its answer is to bolt extra heads onto the same model, predict several future tokens, arrange the guesses as a tree, then use one verification pass to accept as much as possible. The common mistake is to treat this as just another decoding trick. It is really a product engineering compromise: less model management, no separate draft lifecycle, and speedups that depend on acceptance behavior, batching, and workload shape. If you serve LLMs, this paper is worth an evening because it shows how algorithm design changes when HBM traffic, not FLOPs, is the wall.
What to take away
- Medusa replaces a separate draft model with extra decoding heads attached to the served model.
- Tree attention lets one forward pass verify multiple candidate continuations, not just one linear guess.
- Medusa-1 preserves the frozen backbone, while Medusa-2 trades more training complexity for higher acceptance rates.
Reads with
- Fast Inference from Transformers via Speculative Decoding
this is the clean draft and verify baseline medusa removes the extra model from
- Efficient Memory Management for Large Language Model Serving with PagedAttention
medusa speedups meet the real serving bottleneck of batching and kv cache memory
- Fast Transformer Decoding: One Write-Head is All You Need
both trade model structure for lower decode time when memory bandwidth is the wall