Source-linked AI summary
Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding
Zachary Ankner, Rishab Parthasarathy, Aniruddha Nrusimha, Christopher Rinard, Jonathan Ragan-Kelley, William Brandon
TL;DR
Autoregressive LLM inference is memory-bandwidth-bound because tokens are generated sequentially, motivating speculative decoding with lightweight draft heads. Hydra replaces sequentially independent heads with sequentially dependent heads that use earlier candidate tokens, and Hydra++ further tunes their objectives and architectures. Hydra++ improves throughput by up to 1.31× over Medusa decoding and 2.70× over standard autoregressive decoding.
Problem
Sequential autoregressive decoding is memory-bandwidth-bound, while existing draft heads predict candidate tokens independently of earlier candidate-continuation tokens.
Method
Hydra makes each draft-head output depend on the candidate continuation so far, and Hydra++ extends the design with tuned training objectives and architectures.
Results
Hydra++ improves decoding throughput by up to 1.31× over Medusa decoding and 2.70× over standard autoregressive decoding.
Takeaways & Limitations
Sequential dependence is a simple drop-in intervention that improves draft-head speculation and end-to-end decoding throughput.
Abstract
from arXiv · showhide
To combat the memory bandwidth-bound nature of autoregressive LLM inference, previous research has proposed the speculative decoding frame-work. To perform speculative decoding, a small draft model proposes candidate continuations of the input sequence that are then verified in parallel by the base model. One way to specify the draft model, as used in the recent Medusa decoding framework, is as a collection of lightweight heads, called draft heads, that operate on the base model's hidden states. To date, all existing draft heads have been sequentially independent, meaning that they speculate tokens in the candidate continuation independently of any preceding tokens in the candidate continuation. In this work, we propose Hydra heads: a sequentially-dependent drop-in replacement for standard draft heads that significantly improves the accuracy of draft head speculation. We further explore the design space of Hydra head training objectives and architectures, and propose a carefully tuned Hydra head recipe, which we call Hydra++, that improves decoding throughput by up to 1.31x and 2.70x compared to Medusa decoding and autoregressive de-coding respectively. Overall, Hydra heads are a simple and well-motivated intervention on standard draft heads that significantly improve the end-to-end speed of draft head-based speculative decoding. We make our code publicly available at https://github.com/zankner/Hydra.
1 Introduction
Autoregressive LLM inference is limited by sequential, memory-bandwidth-bound decoding, motivating speculative decoding with lightweight draft heads. Hydra introduces sequential dependence between draft-head predictions, while Hydra++ further improves throughput and supports batched and non-greedy settings.
- Motivation: Speculative decoding reduces weight movement per generated token by having a draft model propose multiple tokens that the base model verifies in one forward pass.This targets the memory-bandwidth bottleneck caused by sequential autoregressive generation.
- Hydra heads: Hydra heads use earlier candidate-continuation tokens as inputs, addressing the sequential independence of standard draft heads and increasing average acceptance length by up to 0.46 tokens.Standard heads depend only on hidden states from previously verified tokens; Hydra makes each output a function of the candidate continuation so far.
- Hydra heads: Hydra decoding improves end-to-end throughput by up to 1.10× over Medusa by introducing sequential dependence into draft-head speculation.
- Hydra++: Hydra++ achieves up to 1.31× and 2.70× higher throughput than Medusa decoding and standard autoregressive decoding, respectively.The recipe combines deeper draft-head MLPs, teacher distillation, and an additional transformer decoder layer.
- Alternative settings: Hydra decoding achieves better throughput than Medusa at every evaluated batch size.
- Alternative settings: With typical acceptance sampling, Hydra++ matches base-model sampling quality while retaining speculation's throughput benefits.Typical acceptance sampling is described as a non-distribution-preserving verification criterion.
2 Background
Speculative decoding uses a cheap draft model to propose candidate continuations and a base model to verify them in parallel. Background approaches include tree-structured candidates and lightweight draft heads, with Medusa combining both techniques.
- Speculative decoding: Speculative decoding has a draft model generate candidate continuations, then uses one base-model forward pass to verify them in parallel.Accepted tokens are appended while rejected tokens are discarded.
- Speculative decoding: More accurate draft predictions can improve acceptance under common verification criteria, including rejection resampling and greedy acceptance.
- Tree decoding: Tree decoding represents candidate tokens as tree nodes, with each path corresponding to a possible continuation and children representing possible next tokens.
- Tree decoding: The base model evaluates all tree nodes in one forward pass by packing them together and using an attention mask that restricts each token to its tree parents.
- Lightweight heads and Medusa: Lightweight draft heads are small MLPs operating on the base model's hidden state, with separate heads predicting tokens at different future offsets independently.Medusa uses tree-based speculative decoding with a collection of draft heads, typically implemented as single-layer residual MLPs.
- Lightweight heads and Medusa: Hydra's figure contrasts standard heads, which use only the latest processed-token hidden state, with Hydra heads, which additionally use earlier candidate-continuation tokens.
3 Hydra Heads
Hydra heads replace sequentially independent draft heads with heads that condition on previously sampled candidate tokens. Hydra++ further improves this design through changes to head depth, training objective, and context encoding.
- Hydra Heads: Standard draft heads independently predict future tokens from a hidden state that reflects only the already decoded sequence.They are unaware of earlier tokens in the current candidate continuation.
- Hydra Heads: Hydra heads condition each future-token prediction on the base model state and tokens sampled by previous Hydra heads, unlike standard draft heads.This introduces sequential dependence into candidate generation.
- Hydra Heads: The basic Hydra head is a single-hidden-layer MLP that concatenates the base hidden state with embeddings of preceding candidate tokens.The architecture can support other implementations of fHydra,i.
- Hydra++: Hydra++ combines four-layer MLP heads, self-distillation toward the base model distribution, and prefix attention for broader context conditioning.These changes are presented as beneficial draft-head improvements beyond sequential dependence.
4 Discovering performant decoding trees
The decoding trees are static, data-driven structures selected for the target inference configuration. Proposal trees maximize expected acceptance length approximately, after which measured end-to-end throughput selects the final tree.
- Design considerations: Tree topology is configuration-dependent because model, draft model, batch size, and hardware can change relative performance.This motivates selecting trees separately for each desired inference configuration.
- Tree discovery: The method first constructs proposal trees of increasing sizes and then chooses the tree with maximum measured throughput for the target setup.The setup includes factors such as batch size and hardware.
- Tree discovery: Each proposal tree is built greedily by adding the child node that most improves simulated expected acceptance length.The process starts from a one-node tree and repeats until the desired sequence of sizes is obtained.
- Implementation: In practice, the search uses a maximum tree size of 100 and a 100-question Alpaca subset for simulated statistics.The discovered trees are detailed by decoding strategy and batch size in Appendix B.
5 Shared training and evaluation details
The experiments use frozen Vicuna base models, ShareGPT training data, and MT-Bench evaluation. Throughput and average acceptance length are the main metrics under greedy verification, with separate quality evaluation for non-greedy settings.
- Overview: The shared procedure covers the common training and evaluation setup for the experiments.
- Models and training: Experiments evaluate 7B, 13B, and 33B Vicuna models, with draft heads trained while the base-model weights remain frozen.Training uses the ShareGPT multi-turn conversation dataset.
- Evaluation: MT-Bench is the evaluation benchmark, with throughput measured in tokens per second and acceptance length measured in tokens generated per decoding step.7B and 13B experiments use one A100-40GB GPU, while 33B experiments use one A100-80GB GPU.
- Evaluation: Unless otherwise specified, experiments use greedy verification, so generation quality matches the base model and evaluation focuses on speed and acceptance length.Non-greedy experiments separately assess generation quality.
6 Results
Hydra and Hydra++ improve throughput and acceptance length over Medusa and autoregressive decoding, including batched inference. Under typical acceptance sampling, Hydra++ matches base-model quality while retaining high acceptance length.
- Batch-size-1 throughput: Hydra++ improves throughput over autoregressive decoding by 2.70×, 2.50×, and 2.53× for 7B, 13B, and 33B models, and over Medusa by 1.27×, 1.27×, and 1.31×.Hydra also improves over autoregressive decoding by 2.36×, 2.17×, and 2.15× and over Medusa by 1.11×, 1.10×, and 1.11× for those model sizes.
- Batch-size-1 throughput: Across evaluated model sizes, Hydra has higher average acceptance length than Medusa, while Hydra++ exceeds Hydra and both improvements increase decoding throughput.The results support sequential dependence as an intervention that improves draft prediction accuracy and decoding speed.
- Batched inference: All speculative methods outperform autoregressive decoding at examined batch sizes, but their relative gains decrease as batch size increases.For Hydra++, the improvement over standard decoding falls from 2.70× at batch size 1 to 1.63× at batch size 8.
- Batched inference: For every examined batch size, Hydra decoding remains an improvement over both Medusa and standard decoding despite smaller gains at larger batches.The batched evaluation uses a 7B base model and batch sizes {1,2,4,8}.
- Typical acceptance sampling: Typical acceptance is evaluated by varying posterior threshold ϵ with α = √ϵ and fixed sampling temperature τ = 0.7.The evaluation reports average LLM-as-a-judge scores for MT-Bench Writing and Roleplay categories.
- Typical acceptance sampling: With typical acceptance sampling, Hydra and Hydra++ retain higher average acceptance length than Medusa, and Hydra++ matches base-model generation quality at ϵ = 0.15.Medusa and Hydra do not match random sampling from the base model at the examined thresholds.
7 Related Work
The paper situates Hydra within work on accelerating LLM inference, speculative decoding, draft-model design, memory reduction, and batched inference.
- Speculative decoding research includes retrieval-based proposals, Jacobi iteration, tree verification, hierarchical draft models, online draft-model training, and distillation-based alignment.
- Related work also reduces LLM memory impact through weight quantization, feature pruning, and multi-query or grouped-query attention for KV-cache compression.
- Other studies improve inference throughput by increasing batch size through better scheduling and shared-resource management.
8 Conclusion
The paper concludes that sequential dependence improves draft-head speculation, while Hydra++ combines architectural and training changes for larger throughput gains than Medusa and autoregressive decoding.
- Hydra uses candidate-continuation input embeddings to replace sequentially independent draft heads, achieving up to 1.11× higher throughput than Medusa decoding.
- Hydra++ combines deeper MLPs, teacher distillation, and an extra transformer decoder layer, reaching up to 1.31× and 2.70× higher throughput than Medusa and autoregressive decoding.
- Hydra++ retains benefits in batched inference, and typical acceptance sampling matches non-greedy base-model generation quality without reducing acceptance length.
A.1.1 Results
The results show that Hydra’s training and architecture choices improve speculation and throughput, while tree size must be tuned by decoding strategy and batch size.
- Training interventions: Teacher-loss training improves Vicuna 7B throughput by 1.04× over vanilla Hydra heads, whereas adding input noise degrades acceptance length and decoding speed.
- Hydra-specific prefix attention: PrefixMLP improves average acceptance length by 1.12× and average decoding throughput by 1.08× over MLP-only Hydra heads on MT-Bench.
- Optimal decoding trees: As batch size increases, the tree size maximizing throughput decreases for Medusa, Hydra, and Hydra++ decoding.
- EAGLE comparison: EAGLE achieves higher average acceptance length than Hydra++, but both methods deliver comparable decoding throughput because EAGLE incurs greater draft-model overhead.
- Overhead analysis: Hydra decoding has greater overhead than Medusa, yet still improves end-to-end throughput despite the added prefix-attention and draft-head costs.
- Benchmark results: Hydra++ achieves significantly better throughput than Medusa across all SpecBench task categories.
D Analysis of Hydra Head overheads
Hydra’s sequential dependence and prefix attention add draft-side computation, but the resulting overhead can still accompany end-to-end throughput gains.
- Hydra heads incur increasing first-layer input width for later speculative positions because they consume embeddings of preceding speculated tokens.
- The overhead analysis separates the costs of Hydra’s additional decoder layer from those caused by sequential dependence.
E Evaluation on SpecBench
Hydra++ is evaluated on SpecBench, a broader speculative-decoding suite spanning multiple task types. It significantly outperforms Medusa across all tasks, with gains similar to those observed on MT-Bench.
- Hydra++ outperforms Medusa across all SpecBench tasks, averaging 1.24× higher throughput.SpecBench covers multi-turn chat, translation, summarization, question answering, mathematics, and retrieval-augmented generation.
- The 1.24× average improvement over Medusa on SpecBench closely matches the 1.27× MT-Bench speedup, indicating gains across varied tasks.
- Summary and RAG tasks show the smallest throughput improvements for both Medusa and Hydra++.