EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees

TL;DR

EAGLE-2 enhances speculative sampling for large language models by introducing a context-aware dynamic draft tree, significantly speeding up inference while maintaining output quality.

Contribution

It proposes a novel context-aware dynamic draft tree technique for speculative sampling, improving inference speed over previous static methods.

Findings

Achieves 3.05x-4.26x speedup ratios

20%-40% faster than EAGLE-1

Maintains unchanged output distribution

Abstract

Inference with modern Large Language Models (LLMs) is expensive and time-consuming, and speculative sampling has proven to be an effective solution. Most speculative sampling methods such as EAGLE use a static draft tree, implicitly assuming that the acceptance rate of draft tokens depends only on their position. Interestingly, we found that the acceptance rate of draft tokens is also context-dependent. In this paper, building upon EAGLE, we propose EAGLE-2, which introduces a new technique of context-aware dynamic draft tree into drafting modeling. This improvement leverages the fact that the draft model of EAGLE is well-calibrated: the confidence scores from the draft model approximate acceptance rates with small errors. We conducted extensive evaluations on three series of LLMs and six tasks, with EAGLE-2 achieving speedup ratios 3.05x-4.26x, which is 20%-40% faster than EAGLE-1. EAGLE-2 also ensures that the distribution of the generated text remains unchanged, making it a lossless acceleration algorithm.