AI Intelligence // signal over noise
← back to feed
Cohere 8/10 signal

Hardware-aware Dynamic Speculative Decoding

infrastructureresearch
Summary
Cohere has developed and contributed Hardware-aware Dynamic Speculative Decoding to the open-source vLLM framework. The method dynamically adjusts the number of speculative draft tokens based on batch size, improving inference speed by up to 23% over fixed-token speculative decoding on dense models at high batch sizes.
Context
Speculative decoding is an inference acceleration technique where a small, fast "draft" model proposes several candidate tokens, which a larger "target" model then verifies in a single pass. This exploits the fact that modern GPUs are often memory-bandwidth-bound at low batch sizes, leaving compute units idle. Standard implementations use a fixed number of draft tokens (K), which is suboptimal. As batch sizes increase, inference can become compute-bound, and a high, fixed K can actually degrade throughput. Cohere's work addresses this by making K adaptive, optimizing performance across different serving loads by responding to the shifting bottleneck between memory bandwidth and compute.
Details

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):

ComparisonBatch SizePerformance Gain of DSD
vs. Fixed-K Speculative Decoding128-256~23% faster
vs. Vanilla Inference1287.5% faster
vs. Vanilla Inference2561.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_size and K.
What's new
The novel contribution is making the number of speculative draft tokens (K) in speculative decoding dynamic and hardware-aware. Instead of using a static, pre-set K, this method adapts K at runtime based on the batch size, using an offline-profiled lookup table to maximize throughput as the system bottleneck shifts between memory-bound and compute-bound.
The take

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.