Source-linked AI summary
Accelerating LLM Inference with Staged Speculative Decoding
Benjamin Spector, Chris Re
TL;DR
Local small-batch LLM inference has low arithmetic intensity and creates challenges for latency, personalization, and privacy. Staged speculative decoding restructures speculative batches as trees and speculatively decodes the draft model, achieving an average 3.16x speedup over standard single-batch inference without degrading model quality.
Problem
Local small-batch LLM inference suffers from poor compute utilization and creates challenges for latency, personalization, and privacy.
Method
The method restructures speculative batches as trees and adds speculative decoding for the draft model in a staged oracle–draft–draft2 process.
Results
3.16x average speedup over standard single-batch inference was achieved, while speculative decoding does not degrade model quality.
Takeaways & Limitations
The techniques improve speculative decoding performance in both deterministic and sampling-based decoding while preserving model quality.
Takeaways & Limitations
Performance benefits vary substantially, ranging from 2x to 10x depending largely on the denseness or sparseness of difficult content.
Abstract
from arXiv · showhide
Recent advances with large language models (LLM) illustrate their diverse capabilities. We propose a novel algorithm, staged speculative decoding, to accelerate LLM inference in small-batch, on-device scenarios. We address the low arithmetic intensity of small-batch inference by improving upon previous work in speculative decoding. First, we restructure the speculative batch as a tree, which reduces generation costs and increases the expected tokens per batch. Second, we add a second stage of speculative decoding. Taken together, we reduce single-batch decoding latency by 3.16x with a 762M parameter GPT-2-L model while perfectly preserving output quality.
1. Introduction
Local small-batch LLM inference is important for latency, personalization, and privacy but suffers from low arithmetic intensity. The paper extends speculative decoding with tree-structured batches and an additional speculation stage, reporting improved performance while preserving output quality.
- Motivation: Local inference improves responsiveness, personalization, and privacy while addressing poor compute utilization from low arithmetic intensity.The paper also frames efficient local execution as supporting AI democratization.
- Prior work: Speculative decoding uses a fast draft model to anticipate tokens for the oracle model, but long correct sequences become exponentially unlikely.This causes performance gains to saturate as more sequential tokens are speculated.
- Contributions: The paper restructures speculative batches as trees to create larger, higher-quality batches more quickly.The tree is intended to increase the expected number of useful speculative tokens.
- Contributions: A second speculative stage decodes the draft model itself to further improve performance.This extends speculation beyond the oracle-model query.
- Results: The techniques improve speculative decoding in both deterministic and sampling-based decoding.The reported improvement is described as preserving model output quality.
2. Background
Small-batch autoregressive decoding is expensive because sequential token generation repeatedly streams model weights at very low arithmetic intensity. Speculative decoding trades additional computation for reduced bandwidth use, accelerating this bottleneck without degrading model quality.
- Inference background: Decoding repeatedly selects a token, feeds it back into the model, and produces the next logits until generation ends.Unlike prompt processing, this phase is inherently sequential.
- Inference background: Low arithmetic intensity makes small-batch decoding especially expensive because the entire model’s weights are streamed for each generated token.The passage defines arithmetic intensity as FLOP of compute divided by bytes of memory bandwidth.
- GPU performance: Batch-size-1 GPT-2 Large inference requires approximately 1.4 GFLOP yet reaches only 150 tokens/second and 0.13% compute utilization on an RTX 4090.The example illustrates the GPU utilization challenge in small-batch inference.
- Speculative decoding: Speculative decoding feeds several draft-model predictions to the oracle as one batch, saving memory bandwidth when predictions are accepted.Rejected predictions are discarded and decoding reverts naturally to token-by-token execution.
- Speculative decoding: Speculative decoding is attractive because it preserves model quality and offers gains generally orthogonal to other methods.Its performance comes from converting sequential execution into parallel execution.
3. Methods
Staged speculative decoding combines tree-structured speculative batches with speculation over the draft model. The design targets the diminishing returns and draft-generation costs of standard speculative decoding while reducing memory-bandwidth demands.
- Method overview: Staged speculative decoding combines tree-structured batches with additional speculation stages.The paper names this combination “staged speculative decoding.”
- Motivation: Single-sequence speculation scales poorly because agreement between draft and oracle models over long sequences is exponentially unlikely.This produces rapidly diminishing returns as arithmetic intensity increases.
- Tree-structured batches: The prediction tree increases expected true tokens per batch, leaf nodes, and draft-model parallelism.These benefits arise from dynamically building multiple possible sequences.
- Tree-structured batches: A wider tree can provide more batch predictions without running the draft model at every leaf.The draft model runs at internal nodes, while leaf-node predictions are obtained through the tree structure.
- Implementation: Tree construction controls positional embeddings and causal masking while separately storing the batch KV cache.The implementation partitions attention between the KV cache and self-attention within the batch.
- Staged speculation: Because large speculative batches can make draft generation costlier than oracle execution, the method adds speculative decoding to the draft model.The resulting staged method consists of oracle, draft, and draft2 models with tree-structured batches.
4. Results
Across 164 HumanEval prompts, staged speculative decoding reduced bandwidth and improved throughput over both standard decoding and standard speculative decoding in deterministic and topk settings, though gains varied by prompt.
- Evaluation setup: The evaluation used a 762M-parameter GPT-2-Large oracle, a 40M GPT-2 draft model, and a Katz backoff trigram draft2 model on a quiesced RTX 4090.The methods were compared on 164 HumanEval prompts using deterministic and topk sampling.
- Bandwidth: Staged speculative decoding used substantially less memory bandwidth than both non-speculative and standard speculative decoding.Bandwidth was measured to validate the approach’s intended reduction in memory traffic.
- Throughput: 3.16x average performance boost over the reference implementation and 1.36x over standard speculative sampling was achieved with deterministic sampling.Profiling attributed 35% overhead to the Python infrastructure.
- Per-problem variation: Figure 2 compares relative performance across HumanEval problems for greedy and Topk decoding, with problem indices sorted by staged speculative performance.The distributions show that performance varies across prompts rather than remaining uniform.
- Throughput: 1.98x average performance boost over baseline and 1.36x over standard speculative sampling was achieved with topk sampling at k = 50 and T = 1.Both speculative methods were degraded by stochastic rejection of tokens, but staged speculation retained its lead.
- Token origins: Figure 3 shows easiest tokens, such as whitespace, originating from the N-gram model, while more difficult tokens are generated by larger models and all tokens are checked by the oracle.The displayed prompt achieved approximately 2.5x performance over baseline.
- Performance variation: Performance benefits ranged from 2x to 10x depending substantially on the density or sparseness of difficult content.Highly indented Python code benefited more from the N-gram models than unindented code.
5. Conclusions
The paper combines tree-structured speculative batches with a second speculation stage for the draft model, achieving a 3.16x average speedup over standard single-batch inference.
- Conclusions: The method restructures the oracle batch as a tree to reduce generation cost and increase expected tokens per batch.It also adds a second speculation stage to accelerate draft-model decoding.
- Conclusions: 3.16x average speedup over standard single-batch inference was achieved by combining both improvements.The reported result summarizes the overall staged speculative decoding approach.