Hardware-aware Dynamic Speculative Decoding
Methodology:
- The core idea is to make the number of speculative draft tokens (K) dynamic, adjusting it based on the runtime batch size to maximize performance.
- The system tracks two key metrics: Acceptance Length (AL), the number of drafted tokens accepted by the target model, and Inter-Token Latency (ITL), the total time for drafting and verification.
- A metric called goodput is defined as AL divided by ITL.
- The system performs offline profiling to measure these values across a range of batch sizes and builds a lookup table that maps each batch size to the optimal K value that maximizes goodput. This table is then used at runtime.
Performance on Dense Models (Command A):
| Comparison | Batch Size | Performance Gain of DSD |
|---|---|---|
| vs. Fixed-K Speculative Decoding | 128-256 | ~23% faster |
| vs. Vanilla Inference | 128 | 7.5% faster |
| vs. Vanilla Inference | 256 | 1.82% faster |
- The benchmark dataset used was MT-Bench.
- Fixed-K speculative decoding was observed to regress (become slower than vanilla inference) at these high batch sizes, while the dynamic method provided a speedup.
- For dense models, the optimal K value decreases monotonically as the batch size increases.
Performance on MoE Models (Command A+):
- On the Mixture-of-Experts model Command A+, Dynamic Speculative Decoding (DSD) and fixed-K decoding delivered similar speedups.
- This was because the system selected the same optimal K values across most batch size ranges for this specific model.
- The relationship between optimal K and batch size for MoE models is non-monotonic due to the added complexity of loading different experts.
vLLM Integration:
- The work was contributed upstream to the vLLM open-source inference framework.
- Key engineering challenges included solving asynchronous scheduling to handle a variable number of draft tokens across different timesteps.
- It also required extending full CUDA graph support to capture multiple valid combinations of
batch_sizeand K.
This is a sharp piece of production inference engineering that addresses a real-world problem: static configurations don't work well in dynamic serving environments. The core insight is that the optimal number of speculative tokens changes as the system's bottleneck shifts from memory to compute with increasing batch sizes. By profiling this relationship and creating a simple runtime lookup, Cohere gets significant gains where fixed-K methods fail or even regress. Contributing this to vLLM is a major win for the community, making a sophisticated optimization accessible. It's a reminder that huge performance improvements are still found in the software stack, not just in model architecture. The less dramatic results on MoE models also highlight that optimizations for dense transformers don't always translate directly to more complex architectures.
Don't read this site daily. Get it in your inbox.
The daily brief and Sunday deep dive — distilled, scored, and opinionated. For builders only.