Make Every Draft Count: Hidden State based Speculative Decoding

TL;DR

This paper introduces a novel speculative decoding system that reuses discarded draft hidden states to improve inference efficiency in large language models, achieving up to 3.3x speedup.

Contribution

It proposes a hidden state-based speculative decoding method that transforms discarded drafts into reusable tokens, enhancing efficiency over traditional token-based approaches.

Findings

Achieves up to 3.3x speedup over standard speculative decoding.

Introduces a draft model architecture based on auto-regressive hidden states.

Designs an efficient token injection mechanism for high-quality draft reuse.

Abstract

Speculative decoding has emerged as a pivotal technique to accelerate LLM inference by employing a lightweight draft model to generate candidate tokens that are subsequently verified by the target model in parallel. However, while this paradigm successfully increases the arithmetic intensity of memory-bound inference, it causes significant compute inefficiency: the majority of draft tokens fail verification and are discarded, resulting in waste of computation. Motivated by the goal of recollecting this wasted computation, we propose a novel system that transforms discarded drafts into reusable tokens. Our key insight is to perform auto-regressive prediction at the hidden states level and postpone the integrating token information after the hidden states generation, so the draft hidden states are not contaminated by incorrect tokens, enabling hidden state reuse. To implement such a system, first we introduce a draft model architecture based on auto-regressive hidden states, which preserves richer semantics than token-based drafters to facilitate draft repurposing. Second, we design an efficient token information injection mechanism that leverages our specialized draft model to construct high-quality draft token trees and enables resampling tokens from verification failures. Third, we eliminate the overhead hidden in our design to further maximize hardware utilization. We conducted extensive evaluations against various baselines, demonstrating up to a 3.3x speedup against standard speculative decoding.