AWS shows how speculative decoding on Trainium2 speeds up generation in vLLM

AWS has demonstrated a practical way to make LLM inference on Trainium2 faster and cheaper for workloads where the model generates far more tokens than it receives as input. The technique is speculative decoding. Instead of forcing a large model to emit tokens one at a time, the system adds a small draft model that quickly proposes several upcoming tokens at once, and the main target model verifies them in a single pass. When the guesses match, the service spends fewer expensive sequential steps, cuts inter-token latency and keeps the accelerator busier. This matters most for decode-heavy workloads: writing assistants, coding agents, report generation, templated documents and other tasks with long outputs.

In standard autoregressive generation, each new token is computed separately, so the accelerator keeps reading the KV-cache from memory while doing relatively little useful work per step. As a result, inference is often bound by memory bandwidth rather than raw compute. Speculative decoding targets exactly this bottleneck: the target model runs fewer sequential decode steps, and verifying a batch of tokens makes the workload denser.

The approach does come with requirements. The draft and target models must share the same tokenizer and vocabulary, and ideally belong to the same architecture family, so the small model guesses the main model's continuation more often. The key parameter is the number of speculative tokens. If the window is too small, the gain is barely noticeable; if it is too large, early rejections and extra verification eat up the benefit.

In its test, AWS used Qwen3-32B as the target model and Qwen3-1.7B as the draft model, served through vLLM on a trn2.48xlarge Trn2 instance. For speculative decoding, the team chose fused speculation in NeuronX Distributed Inference, which compiles both models together for better performance. The baseline and speculative configurations were deployed in the same Amazon EKS cluster with everything else held constant: accelerator allocation, tensor parallelism, context length, batch limits and the Neuron image. The only difference was the added draft model and the num_speculative_tokens parameter. Both services were load-tested with llmperf, and TTFT, inter-token latency and end-to-end latency were sent to CloudWatch for comparison.

AWS also tested the smaller Qwen3-0.6B separately, but its acceptance rate came in roughly 60 percent lower, which was enough to wipe out most of the gain. Across a range of 5 to 15 speculative tokens, the seven-token configuration performed best in these tests, although the authors specifically stress that the optimal value depends heavily on prompt structure.

Prompt structure ultimately decided the outcome. On predictable workloads, such as repetitive text, numeric sequences and simple code, speculative decoding delivered a clear benefit. In these cases, the draft model often guesses what the target model would have produced anyway, so the system skips a substantial share of sequential steps. In the tests, inter-token latency dropped to roughly 15 milliseconds per token, and the end-to-end latency curve consistently stayed below baseline.

Open-ended, less deterministic prompts told a different story. The draft model diverges from the target model more often, tokens get rejected, and the potential gain disappears. For these prompts, inter-token latency held at around 45 milliseconds per token, and the speculative and baseline configurations showed nearly identical overall latency.

TTFT, or time to first token, barely changed, because speculative decoding does not speed up the prefill stage, where the model encodes the input context. The main benefit arrives later, in the decode phase, by reducing the number of expensive sequential steps the target model has to take.

The practical takeaway is simple: speculative decoding on Trainium2 is not a universal speed-up button but a targeted optimization for a specific type of workload. If a product frequently generates structured, predictable output (code, data extraction, templated reports, configs), this mode can directly lower the cost per output token and raise throughput without sacrificing quality. If your traffic is mostly open-ended chat with free-form generation, the effect may be minimal. That's why this setup is worth adopting only after measuring it on your own prompts, choosing a compatible draft model and tuning the speculative-token window to real-world scenarios, rather than relying on benchmarks divorced from the product.