title: Medusa: Simple LLM Inference Acceleration Framework with Multiple
    Decoding Heads
summary: 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.
paper_authors._id: 65add289ac588f2a1c2b22c2
paper_authors.name: Tianle Cai
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-22T11:22:03.769Z
paper_authors._id: 64459357c525660aa20be337
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/noauth/1fryBXRuSvAgMIDkpWBpI.jpeg
paper_authors.fullname: Tianle Cai
paper_authors.type: user
paper_authors.user: tianlecai
paper_authors._id: 65add289ac588f2a1c2b22c3
paper_authors.name: Yuhong Li
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-22T11:22:29.834Z
paper_authors._id: 63687abf46919b9619c157f2
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/1667791517323-noauth.png
paper_authors.fullname: Yuhong Li
paper_authors.type: user
paper_authors.user: yli3521
paper_authors._id: 65add289ac588f2a1c2b22c4
paper_authors.name: Zhengyang Geng
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-22T11:22:37.019Z
paper_authors._id: 626039656089b62d9f332854
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/1650473291256-noauth.jpeg
paper_authors.fullname: Zhengyang Geng
paper_authors.type: user
paper_authors.user: Gsunshine
paper_authors._id: 65add289ac588f2a1c2b22c5
paper_authors.name: Hongwu Peng
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-22T11:22:55.102Z
paper_authors._id: 655a9f12bf604a7732716115
paper_authors.avatarUrl: /avatars/2f5094af95bf3eb236e39dacb909a5d7.svg
paper_authors.fullname: Hongwu Peng
paper_authors.type: user
paper_authors.user: harveyp123
paper_authors._id: 65add289ac588f2a1c2b22c6
paper_authors.name: Jason D. Lee
paper_authors._id: 65add289ac588f2a1c2b22c7
paper_authors.name: Deming Chen
paper_authors._id: 65add289ac588f2a1c2b22c8
paper_authors.name: Tri Dao
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-01-22T11:23:52.871Z
paper_authors._id: 64b8a6b5cf14c2fabe98159b
paper_authors.avatarUrl: /avatars/dbc009451865435bf290791beadc4723.svg
paper_authors.fullname: Tri Dao
paper_authors.type: user
paper_authors.user: tridao
paper_title: Medusa: Simple LLM Inference Acceleration Framework with Multiple
    Decoding Heads
paper_summary: 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.
paper_ai_keywords: auto-regressive decoding memory bandwidth speculative decoding draft model Medusa decoding heads tree-based attention mechanism candidate continuations parallel processing single-step latency Medusa-1 Medusa-2 fine-tuning procedures frozen backbone lossless inference acceleration prediction accuracy self-distillation acceptance scheme generation quality speedup