DynaSpec: Context-aware Dynamic Speculative Sampling for Large-Vocabulary Language Models

TL;DR

DynaSpec addresses the bottleneck of large-vocabulary drafting in speculative decoding by introducing a context-aware, dynamic shortlisting mechanism. A lightweight meta-classifier routes each draft step to a small union of token clusters, reducing the per-token head computation while preserving exact verification over the full vocabulary; a position-aware budget further balances early versus late steps. Theoretical analysis shows dynamic routing dominates static shortlists in retained probability mass, improving expected tokens per verification step. Empirically, DynaSpec yields higher mean accepted length and faster decoding across seven tasks with 128k vocabularies, approaching full-vocabulary performance and achieving up to 2.18x speedups over fixed-shortlist baselines, while remaining plug-in to EAGLE-style pipelines.

Abstract

Speculative decoding has become a standard way to accelerate LLM inference: a small drafter proposes multiple tokens and a large target model verifies them once per speculation length. Recently, scaling of the LLM vocabulary has pushed the number of tokens to grow substantially. While verification over the full vocabulary leaves the target model largely unaffected, the O(|V|d) parameters in the drafter's output head become a latency bottleneck, slowing the entire pipeline. Contemporary methods (e.g., FR-Spec, VocabTrim) restrict the drafter's vocabulary to a fixed top frequent subset of the target model's vocabulary. Although this reduces draft-time compute, it is brittle, since: (i) frequency lists are corpus-dependent and require retuning to generalize, and (ii) static shortlists suppress rare or domain-specific tokens, lowering the expected number of tokens per verification step. We propose DynaSpec, a context-dependent dynamic shortlisting mechanism that is robust, speeds up drafting, and generalizes across diverse tasks. Concretely, we introduce lightweight, coarse-grained meta-classifiers that route contexts to a small number of token clusters; the union of the top-k selected clusters forms the drafter's shortlist, while verification retains the full vocabulary and exactness. The meta-classifier finishes its computation earlier than the drafter's hidden state generation by exploiting parallel execution of draft encoding and meta shortlisting on separate streams. Across standard speculative decoding benchmarks, DynaSpec delivers consistent improvements in mean accepted length, for Llama-3-8B, reaching upto 98.2% of full-vocabulary performance, while fixed-shortlist baselines attain only 84.4%. By leveraging context-dependent selection, DynaSpec achieves up to a 2.18 times increase in generated tokens compared to 1.91 times for fixed-vocabulary approaches.

Figures (3)

• Figure 1: Draft process of DynaSpec. The backbone structure is based on EAGLE where the meta classifier selects the top-k tokens dynamically based on the context.
• Figure 2: Time breakdown of the drafting process of DynaSpec. $T_{\text{emb}}$ is the embedding time, $T_{\text{core}}$ is the encoder layer time, $T_{\text{meta}}$ is the time related to the meta classifier and $T_{\text{index+gemm}}$ is the indexed matmul time.
• Figure 3: Mean accepted length comparison for fixed top-$k$ clusters represented by DynaSpec-F vs. Position-aware top-$k$ clusters represented by DynaSpec-PA across seven datasets using Llama-3-8B model. Mean vocab size represents the mean shortlisted vocab size.