Elamur Research
☰
home / papers / 2024 / Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

Written for the Research library (claude-opus-5-5) from the paper text. Original — arxiv.org/pdf/2401.10774

Summary

Generating text with an LLM is slow because tokens come out one at a time. Each step has to read all of the model's weights from HBM, so the accelerator spends most of its time waiting on memory rather than computing. Speculative decoding helps, but it needs a separate draft model, which is costly to obtain and awkward to deploy in distributed systems. Medusa removes the draft model. It adds a few lightweight extra heads to the final hidden state of the existing LLM, and each head guesses a token further ahead. The top guesses from all heads are combined into a tree of candidate continuations. A special attention mask lets the model check every branch in a single forward pass and keep the longest accepted prefix. There are two training modes: Medusa-1 trains only the heads on a frozen backbone, and Medusa-2 trains the heads and the backbone together with a recipe meant to protect output quality. For models without usable training data, the authors generate it through self-distillation. They also propose a "typical acceptance" rule as a faster alternative to rejection sampling. The result is a 2.3–2.8× speedup on chat models without a noticeable loss in output quality.

Key points

  • Models tested: Vicuna-7B, Vicuna-13B, Vicuna-33B and Zephyr-7B. Experiments focus on batch size 1, which matches local, single-user serving.
  • Medusa-1 (frozen backbone, quality unchanged) gives just over 2.2× speedup. Medusa-2 (joint training) reaches 2.3–2.8×. On Vicuna-7B the figures were 2.18× and 2.83×, while naive joint fine-tuning without the recipe made outputs worse.
  • Ablation of cumulative speedups: heads alone ~1.5×, adding tree attention ~1.9×, an optimized tree layout ~2.2×, and Medusa-2 training ~2.8×.
  • Each head is a single residual SiLU feed-forward layer. Its output projection starts as a copy of the original LM head and its inner weight starts at zero, so at first the heads predict exactly what the base model does. Heads can be trained on one GPU, and QLoRA-style quantization of the frozen backbone is also possible.
  • Self-distillation produces training data for models whose fine-tuning data is private or that went through RLHF. Typical acceptance uses a temperature-dependent threshold to accept plausible candidates and increase accepted length. Published at ICML 2024.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

Fields: AI & machine learning · Computing & GPUs

Cited by

257 papers of our library cite it (of the 140673 whose bibliographies we read); 3 of them have a page here, 254 not yet.

Cite

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2401.10774 · arXiv 2401.10774

@misc{cai2024medusa,
  title = {{Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads}},
  author = {Cai, Tianle and Li, Yuhong and Geng, Zhengyang and Peng, Hongwu and Lee, Jason D. and Chen, Deming and Dao, Tri},
  year = {2024},
  doi = {10.48550/arxiv.2401.10774},
  eprint = {2401.10774},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2401.10774}
}

References

41 references, 7 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F., and Sanghai, S. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245, 2023. find our card
  2. Axolotl. Axolotl. https://github.com/ OpenAccess-AI-Collective/axolotl, 2023. ↗
  3. Dubois, Y., Li, X., Taori, R., Zhang, T., Gulrajani, I., Ba, J., Guestrin, C., Liang, P., and Hashimoto, T. B. Alpacafarm: A simulation framework for methods that learn from human feedback, 2023. find our card
  4. Basu, S., Ramachandran, G. S., Keskar, N. S., and Varshney, L. R. {MIROSTAT}: A {neural} {text} {decoding} {algorithm} {that} {directly} {controls} {perplexity}. In International Conference on Learning Representations, Elfwing, S., Uchibe, E., and Doya, K. Sigmoid-weighted linear units for neural ne… — [doi:10.1016/j.neunet.2017.12.012](https://doi.org/10.1016/j.neunet.2017.12.012) find
  5. Fan, A., Lewis, M., and Dauphin, Y. Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). Association for Computational Linguistics, 2018. doi: 10.18653/v1/p18-1082. find
  6. Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. November 2022. doi: 10.48550/ARXIV.2211.17192. find
  7. Li, X., Zhang, T., Dubois, Y., Taori, R., Gulrajani, I., Guestrin, C., Liang, P., and Hashimoto, T. B. Alpacaeval: An automatic evaluator of instruction-following models. https://github.com/tatsu-lab/ alpaca\_eval, 2023. ↗
  8. Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323, 2022. find
  9. Google. Palm 2 technical report, 2023. URL https://ai.google/static/documents/ palm2techreport.pdf. ↗
  10. Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023. find
  11. Hewitt, J., Manning, C. D., and Liang, P. Truncation sampling as language model desmoothing. October 2022. doi: 10.48550/ARXIV.2210.15191. find
  12. Meister, C., Wiher, G., Pimentel, T., and Cotterell, R. On the probability-quality paradox in language generation. March 2022. doi: 10.48550/ARXIV.2203.17217. find
  13. Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022. find
  14. Meister, C., Pimentel, T., Wiher, G., and Cotterell, R. Locally typical sampling. Transactions of the Association for Computational Linguistics, 11:102–121, 2023. find
  15. Miao, X., Oliaro, G., Zhang, Z., Cheng, X., Wang, Z., Wong, R. Y. Y., Chen, Z., Arfeen, D., Abhyankar, R., and Jia, Z. Specinfer: Accelerating generative llm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023. find
  16. Holtzman, A., Buys, J., Du, L., Forbes, M., and Choi, Y. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum? id=rygGQyrFvH. ↗
  17. NVIDIA. Nvidia a100 tensor core gpu. OpenAI. Gpt-4 technical report, 2023. find our card
  18. Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., and Chen, W. Lora: Low-rank adaptation of large language models. ICLR, 2021. find
  19. Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.02155, 2022. find
  20. Joao Gante. Assisted generation: a new direction toward low-latency text generation, 2023. URL https://huggingface.co/blog/ assisted-generation. ↗
  21. Pan, J. Tiny vicuna 1b. https://huggingface.co/ Jiayi-Pan/Tiny-Vicuna-1B, 2023. ↗
  22. Kim, S., Hooper, C., Gholami, A., Dong, Z., Li, X., Shen, S., Mahoney, M. W., and Keutzer, K. Squeezellm: Dense-and-sparse quantization. arXiv preprint arXiv:2306.07629, 2023. find our card
  23. Pillutla, K., Swayamdipta, S., Zellers, R., Thickstun, J., Welleck, S., Choi, Y., and Harchaoui, Z. MAUVE: Measuring the gap between neural text and human text using divergence frontiers. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing… find our card
  24. Kim, Y. and Rush, A. M. Sequence-level knowledge distillation. EMNLP, 2016. find our card
  25. Kumar, A., Raghunathan, A., Jones, R., Ma, T., and Liang, P. Fine-tuning can distort pretrained features and underperform out-of-distribution. International Conference on Learning Representations, 2022. find
  26. Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference. November 2022. doi: 10.48550/ARXIV.2211.05102. find
  27. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023. find
  28. ShareGPT. ShareGPT. https://huggingface. co/datasets/Aeala/ShareGPT\_Vicuna\_ unfiltered, 2023. find
  29. Shazeer, N. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019. find
  30. Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. P., Zhang, H., Gonzalez, J. E., and Stoica, I. Judging llm-as-a-judge with mt-bench and chatbot arena, 2023. find
  31. Spector, B. and Re, C. Accelerating llm inference with staged speculative decoding. arXiv preprint arXiv:2308.04623, 2023. — [arXiv:2308.04623](https://arxiv.org/abs/2308.04623) find
  32. Stern, M., Shazeer, N. M., and Uszkoreit, J. Blockwise parallel decoding for deep autoregressive models. Neural Information Processing Systems, 2018. find
  33. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. find
  34. Tunstall, L., Beeching, E., Lambert, N., Rajani, N., Rasul, K., Belkada, Y., Huang, S., von Werra, L., Fourrier, C., Habib, N., Sarrazin, N., Sanseviero, O., Rush, A. M., and Wolf, T. Zephyr: Direct distillation of lm alignment, 2023. find
  35. Xia, H., Ge, T., Chen, S.-Q., Wei, F., and Sui, Z. Speculative decoding: Lossless speedup of autoregressive translation, 2023. URL https://openreview.net/forum? id=H-VlwsYvVi. ↗
  36. Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087–38099. PMLR, 2023a. find
  37. Xiao, Y., Wu, L., Guo, J., Li, J., Zhang, M., Qin, T., and Liu, T.-y. A survey on non-autoregressive generation for neural machine translation and beyond. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023b. find
  38. Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in Neural Information Processing Systems, 34:28877–28888, 2021. find
  39. Zhang, P., Zeng, G., Wang, T., and Lu, W. Tinyllama: An open-source small language model, 2024. find our card
  40. Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022. find
  41. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., et al. H 2 o: Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint arXiv:2306.14048, 2023. find