Medusa: Simple LLM Inference Acceleration Framework with Multiple
  Decoding Heads
Authors: name: Tianle Cai; name: Yuhong Li; name: Zhengyang Geng; name: Hongwu Peng; name: Jason D. Lee; name: Deming Chen; name: Tri Dao
Published: 2024-01-19 15:48:40

Abstract: The inference process in Large Language Models (LLMs) is often limited due to
the absence of parallelism in the auto-regressive decoding process, resulting
in most operations being restricted by the memory bandwidth of accelerators.
While methods such as speculative decoding have been suggested to address this
issue, their implementation is impeded by the challenges associated with
acquiring and maintaining a separate draft model. In this paper, we present
Medusa, an efficient method that augments LLM inference by adding extra
decoding heads to predict multiple subsequent tokens in parallel. Using a
tree-based attention mechanism, Medusa constructs multiple candidate
continuations and verifies them simultaneously in each decoding step. By
leveraging parallel processing, Medusa introduces only minimal overhead in
terms of single-step latency while substantially reducing the number of
decoding steps required.
  We present two levels of fine-tuning procedures for Medusa to meet the needs
of different use cases: Medusa-1: Medusa is directly fine-tuned on top of a
frozen backbone LLM, enabling lossless inference acceleration. Medusa-2: Medusa
is fine-tuned together with the backbone LLM, enabling better prediction
accuracy of Medusa heads and higher speedup but needing a special training
recipe that preserves the backbone model's capabilities.
  Moreover, we propose several extensions that improve or expand the utility of
Medusa, including a self-distillation to handle situations where no training
data is available and a typical acceptance scheme to boost the acceptance rate
while maintaining generation quality. We evaluate Medusa on models of various
sizes and training procedures. Our experiments demonstrate that Medusa-1 can
achieve over 2.2x speedup without compromising generation quality, while
Medusa-2 further improves the speedup to 2.3-3.6x.