From AllReduce to EAGLE-3: Building a Speculative Decoding Pipeline for Kimi-K2.5 on Blackwell

Introduction

In the high-stakes world of large-scale ML inference optimization, the gap between a bottleneck and a breakthrough is often bridged by systematic methodology. This article chronicles a pivotal arc within a multi-week coding session dedicated to deploying and optimizing the Kimi-K2.5 INT4 model—a 1-trillion-parameter Mixture-of-Experts (MoE) reasoning model running on 8× NVIDIA RTX PRO 6000 Blackwell GPUs connected only by PCIe. The arc spans from a deep profiling revelation through exhaustive research, a decisive experiment, and ultimately the construction of a complete EAGLE-3 training pipeline.

The central narrative is one of intellectual honesty: the assistant and user confronted a hard hardware constraint (AllReduce consuming 51.5% of decode time), systematically evaluated five speculative decoding options, empirically tested the simplest (n-gram speculation), watched it fail exactly as predicted, and then pivoted to build the most promising path—a custom EAGLE-3 training pipeline. Along the way, they navigated reasoning model benchmarking pitfalls, API incompatibilities between the speculators library and vLLM 0.16, and the unique challenges of extracting hidden states from a multimodal-wrapped MoE architecture.

The Bottleneck That Changed Everything

The story begins in the preceding segment, where an exhaustive profiling campaign using PyTorch's torch.profiler across 8 GPU ranks revealed a stark truth: AllReduce communication consumed 51.5% of decode time ([msg 2471]). For the Kimi-K2.5 INT4 model—a 1T-parameter MoE with 61 layers, 384 routed experts (top-8), and a 163,840-token vocabulary—every single decode step required 127 NCCL AllReduce synchronization barriers traversing PCIe Gen5 links. The measured single-stream throughput was 82.5 tok/s with a 12.1ms time-per-output-token (TPOT), but the GPUs were spending more time talking to each other than computing.

The root cause was unambiguous: the eight RTX PRO 6000 Blackwell GPUs had no NVLink interconnect. All inter-GPU communication traveled over PCIe, and for a model of this scale using Tensor Parallelism (TP=8), the communication tax was crippling. The assistant had already exhausted the obvious optimization paths. Expert Parallelism (EP) was ruled out because PCIe all-to-all lacks multicast capability—measured latency jumped from 19μs at 1 token/GPU to 1,007μs at 16 tokens/GPU. AllReduce fusion was tested and broke on SM120 (Blackwell's compute capability 12.0). MSCCLPP and SBO optimizations yielded negligible improvements. The user had explicitly rejected any precision-cutting hacks, stating "I don't want to do any downcasting hacks on all to all bc need the model to have maximum intelligence" ([msg 2461]).

The team was at a genuine impasse. The hardware physics were immutable.

The Pivot to Speculative Decoding

It was at this moment that the user asked a pivotal question: "Can we investigate speculative decoding?" ([msg 2472]). The user provided two reference articles—a Medium primer on speculative decoding and the Baseten blog post about how they achieved 340+ tok/s on Kimi-K2.5 using a custom EAGLE-3 head on B200 GPUs with NVLink. The user's instructions were precise: "Find a draft model or think what it would take to train one. Look at vllm / sglang support. Use parallel agents for deep research."

The assistant's response was a masterclass in research orchestration ([msg 2473]). Rather than investigating sequentially, the assistant launched five parallel research agents covering: speculative decoding fundamentals and the two reference articles, vLLM's speculative decoding support, SGLang's speculative decoding support, candidate draft models for Kimi-K2.5, and the feasibility of training a custom draft model. These agents ran concurrently as task tool calls, each executing a multi-round conversation before returning results.

The research returned a comprehensive picture that was sobering in several dimensions ([msg 2477]). First, Kimi-K2.5 has no built-in Multi-Token Prediction heads (num_nextn_predict_layers: 0), unlike DeepSeek V3 which has MTP=1, eliminating the "free" speculation path. Second, the tokenizer is incompatible with DeepSeek's family (vocab_size 163,840 vs 129,280), ruling out all DeepSeek-based draft models. Third, the only viable off-the-shelf draft model was AQ-MedAI/Kimi-K2-Instruct-eagle3, trained for K2 (not K2.5), so acceptance rates would be uncertain.

Most critically, the research surfaced the MoE-Spec paper (published just four days before this message, on February 17, 2026), which identified a fundamental problem for speculative decoding on MoE models: during verification of k draft tokens, the target model runs a forward pass on k+1 tokens simultaneously, and each token routes to 8 of 384 experts. Collectively, k+1 tokens may activate 30-60 unique experts per layer, dramatically increasing memory pressure and communication overhead. The paper warned that speculative decoding could cause slowdowns of up to 1.5x on MoE models—exactly the opposite of the intended effect.

The assistant synthesized these findings into a detailed research report ([msg 2477]) that laid out five concrete options (A through E), ranked by effort and expected impact. Option A was n-gram speculation (zero effort, 5-minute configuration change). Option B was switching to SGLang with the existing K2 EAGLE-3 drafter. Option C was trying the K2 EAGLE-3 drafter on vLLM. Option D was training a custom EAGLE-3 head (the Baseten approach, 1-2 weeks). Option E was the cutting-edge MoE-Spec expert budgeting enhancement with no open-source implementation.

The assistant's recommendation was pragmatic: start with Option A (n-gram, 5 minutes) then Option C (existing K2 drafter), using cheap experiments to gather data before committing to the expensive training path.## The User's Directive: Parallel Execution

The user's response ([msg 2479]) was decisive and strategic: "Try option A, before tho, write down what it would take to train Option D in next-steps-eagle.md; Training - setup on the GPUs we have now, if it will look like it would take days I'll get a B200/B300 nvl8 machine to run the hero run."

This directive embodied a dual-track approach. The user wanted the quick experiment (n-gram) to gather empirical data, but also wanted the full training pipeline documented before the experiment began, so that if n-gram failed (as the research predicted), the team could immediately pivot to the more ambitious path without losing time to pipeline setup. The mention of renting B200/B300 NVL8 machines for the "hero run" showed a clear-eyed understanding of the hardware requirements: the current 8× RTX PRO 6000 setup might be sufficient for pipeline development and testing, but the actual training run might need NVLink-connected GPUs.

The assistant acknowledged this with parallel todos ([msg 2480]): writing the EAGLE-3 training guide and testing n-gram speculation, both marked "in_progress" simultaneously.

The EAGLE-3 Training Plan

The assistant wrote next-steps-eagle.md ([msg 2482]), a comprehensive document covering the entire EAGLE-3 training pipeline. The plan included: a draft model configuration matching the K2 EAGLE-3 architecture with a 32K draft vocabulary; a dataset preparation pipeline using HuggingFace datasets to reformat conversations, apply chat templates, and tokenize inputs; a hidden state extraction procedure using the speculators library's VllmHiddenStatesGenerator to capture intermediate representations from the running Kimi-K2.5 model; a vocabulary mapping script to align the draft model's output distribution with the target model's 163,840-token vocabulary; a training script using the speculators library's Eagle3Trainer with knowledge distillation loss; and a shell orchestrator to sequence the entire pipeline.

The document also addressed the hardware decision criterion explicitly: estimating training time on the current 8× RTX PRO 6000 setup versus on rented B200/B300 NVL8 machines, providing the user with the information needed to make the call about the hero run.

The Baseline That Almost Wasn't

While the training document was being written, the assistant simultaneously prepared for the n-gram experiment. This required establishing a performance baseline before modifying the configuration. The assistant wrote a benchmark script (bench_spec_decode.py) and began collecting measurements ([msg 2484]).

The first attempt failed with a 404 error due to a URL construction bug—a double /v1 in the API path ([msg 2485]). The assistant fixed it immediately ([msg 2486]), but the next run revealed a more subtle problem: every prompt returned 0 tokens ([msg 2487]).

This was the reasoning model trap. Kimi-K2.5 generates a hidden reasoning_content field (its chain-of-thought) before producing visible content tokens. With max_tokens=512, the model consumed the entire token budget on reasoning alone, leaving zero tokens for the visible answer. The benchmark was faithfully reporting what it measured—zero content tokens—but the measurement was meaningless.

The assistant diagnosed the issue correctly ([msg 2488]): "Most prompts returned 0 tokens — this is the reasoning model behavior (all tokens go to reasoning_content, and only the final answer comes as content). The 512 max_tokens is being consumed by reasoning tokens." The fix was twofold: increase max_tokens to 4096 and modify the script to count reasoning tokens alongside content tokens.

The corrected baseline revealed the model's true performance: approximately 74 tok/s average across diverse prompts, with a TPOT of ~13.6ms ([msg 2489]). The coding prompt (a binary search tree implementation) achieved 62.6 tok/s with a long reasoning chain, while the explanatory prompt reached 79.5 tok/s. This 27% performance spread showed that throughput depended sensitively on the nature of the reasoning task.

The N-Gram Gamble

With the baseline established and the training document written, the assistant proceeded to the n-gram experiment. The approach was methodical: create a separate systemd service file rather than modifying the existing one ([msg 2490]), preserving the ability to roll back instantly. The new service, vllm-kimi-k25-int4-ngram.service, duplicated the existing configuration but added the critical flag:

--speculative-config '{"method": "ngram", "num_speculative_tokens": 5, "prompt_lookup_max": 4}'

The choice of 5 speculative tokens with a prompt lookup window of 4 was conservative—enough to potentially see a speedup if n-gram worked, but not so aggressive that verification overhead would be crippling.

Before committing to the 30-minute model reload, the assistant verified the configuration was correctly written ([msg 2491]), then stopped the production service ([msg 2492]). The cleanup sequence was aggressive: systemctl stop, followed by kill -9 on any remaining Python/vLLM processes, followed by kill -9 on any processes holding GPU device files, followed by verification with nvidia-smi showing all four GPUs at 0 MiB used.

The assistant then started the n-gram service ([msg 2493]) and began monitoring the logs ([msg 2494], [msg 2495], [msg 2496]). The model loading process was expected to take 30-90 minutes, and the assistant checked periodically for signs of failure.

The Moment of Truth

When the model finally loaded and the assistant ran the benchmark, the results were unambiguous: n-gram speculation was 9–26% slower than baseline ([msg 2499], [msg 2500]). The acceptance rate of drafted tokens was abysmal—under 30%—because the reasoning model's thinking chains were novel and unpredictable, making n-gram matches nearly nonexistent. The MoE expert activation overhead during verification, exactly as the MoE-Spec paper had predicted, negated any benefit from the speculative tokens.

The assistant's analysis ([msg 2502]) was precise: "The n-gram speculation test showed that the MoE verification overhead (activating all 8 experts per verification step) combined with the AllReduce communication tax made speculation slower than baseline." This was a counterintuitive result: speculative decoding is supposed to make things faster, but on PCIe-bound MoE systems, the verification step itself was so expensive that the benefit disappeared.

The assistant restored the original service ([msg 2501]) and documented the findings. The negative result, while disappointing, was valuable: it validated the research, confirmed the bottleneck analysis, and cleared the path for the more ambitious EAGLE-3 training pipeline.## Building the EAGLE-3 Training Pipeline

With n-gram speculation ruled out, the assistant pivoted to implementing the EAGLE-3 training pipeline ([msg 2505]). The user's directive had been clear: "setup on the GPUs we have now, if it will look like it would take days I'll get a B200/B300 nvl8 machine to run the hero run." The assistant would build and test the pipeline on the existing 8× RTX PRO 6000 hardware, with the understanding that the actual training run might be ported to rented machines.

The pipeline construction was systematic. The assistant first researched the speculators library's architecture ([msg 2509], [msg 2510]), reading source code to understand the hidden state extraction mechanism, the training data format, and the vocabulary mapping requirements. The research revealed that the speculators library (v0.3.0) was designed for vLLM ≤0.15, while the installed vLLM was version 0.16—a significant API compatibility concern.

The assistant installed the speculators library ([msg 2512]) and verified the import worked ([msg 2513], [msg 2514]). It then built the pipeline components one by one:

  1. Draft model configuration ([msg 2523]): A configuration matching the K2 EAGLE-3 architecture with a 32K draft vocabulary, designed to be a lightweight ~1B parameter head.
  2. Dataset preparation script ([msg 2524]): Using HuggingFace datasets to load mlabonne/open-perfectblend, reformat conversations, apply the Kimi-K2.5 chat template, and tokenize inputs.
  3. Hidden state extraction script ([msg 2524]): Using the speculators library's VllmHiddenStatesGenerator to capture intermediate representations from the running Kimi-K2.5 model.
  4. Vocabulary mapping script ([msg 2526]): To align the draft model's output distribution with the target model's 163,840-token vocabulary, since the draft model uses a subsampled 32K vocabulary.
  5. Training script ([msg 2527]): Using the speculators library's Eagle3Trainer with knowledge distillation loss.
  6. Shell orchestrator ([msg 2528]): To sequence the entire pipeline from data preparation through training to deployment. The pipeline was tested end-to-end with 10 samples from mlabonne/open-perfectblend ([msg 2531]). Dataset preparation and vocabulary mapping succeeded, but hidden state extraction hit runtime errors due to API mismatches between the speculators library (designed for vLLM ≤0.15) and the installed vLLM 0.16.

Patching Through the API Incompatibilities

The hidden state extraction errors revealed a cascade of incompatibilities. The assistant diagnosed each one systematically:

First, the Kimi-K2.5 model uses a multimodal wrapper architecture. The speculators library expected model.model.layers to access the transformer layers, but Kimi-K2.5's structure is model.language_model.model.layers ([msg 2552]). The assistant patched the library to handle this wrapper.

Second, the SchedulerConfig API had changed between vLLM 0.15 and 0.16. Parameters that were positional in the older version were now keyword-only, and new parameters had been added ([msg 2539], [msg 2540]). The assistant patched the library to match the new API.

Third, the KV cache utility API had changed. The speculators library's hidden state generator relied on internal vLLM APIs for managing the KV cache that had been restructured in v0.16 ([msg 2543], [msg 2544]). The assistant traced the call chain through vLLM's source code to understand the new API surface.

Fourth, the SupportsEagle3 protocol interface had changed. The assistant probed the vLLM source code to understand the new interface requirements ([msg 2545], [msg 2546], [msg 2547], [msg 2548], [msg 2549]), discovering that the DeepSeek V3 model class (which Kimi-K2.5 extends) did not implement the SupportsEagle3 protocol ([msg 2550]).

The assistant applied surgical patches to the speculators library code ([msg 2552], [msg 2554], [msg 2555]), modifying the hidden state generator to handle the multimodal wrapper architecture and the updated vLLM APIs. Each patch was tested by re-running the hidden state extraction script and examining the error messages to determine the next layer of incompatibility.

The Impasse

Despite extensive patching, the hidden state extraction remained blocked by KV cache utility API mismatches that could not be resolved without either upgrading the speculators library to full vLLM 0.16 compatibility (a significant engineering effort) or downgrading vLLM to a version compatible with the library. At the segment's end, the assistant had documented the full training plan and pipeline scripts in next-steps-eagle.md and the eagle3-train/ directory, but the pipeline could not complete its end-to-end run due to the API incompatibilities.

The full training plan and pipeline scripts were documented in next-steps-eagle.md and the eagle3-train/ directory, providing a complete blueprint that could be executed once the compatibility issues were resolved—either by upgrading the speculators library, downgrading vLLM, or porting the pipeline to a framework with better API alignment.

Lessons Learned

This arc of the conversation teaches several important lessons about large-scale ML inference optimization:

1. Empirical testing trumps theoretical prediction. The assistant's research correctly predicted that n-gram speculation would be ineffective for reasoning models, but the user insisted on empirical validation. The 30-minute reload cost was justified by the certainty it provided: the negative result was definitive, not hypothetical.

2. Measurement infrastructure must be adapted to the model. The reasoning model's reasoning_content field nearly invalidated the baseline benchmark. A standard benchmarking approach would have silently produced meaningless data. The assistant's willingness to diagnose and fix the measurement methodology was essential.

3. API compatibility is the hidden tax on innovation. The speculators library's incompatibility with vLLM 0.16 consumed more engineering effort than any other single task in the pipeline. The gap between a library designed for vLLM ≤0.15 and the installed vLLM 0.16 was a chasm of changed interfaces, restructured APIs, and missing protocol implementations.

4. Parallel execution maximizes throughput of the optimization process. By writing the EAGLE-3 training documentation while the n-gram experiment was being prepared, the assistant effectively time-shifted a cognitively demanding task into a period that would otherwise have been idle. This pattern of overlapping workstreams was a recurring theme throughout the session.

5. The MoE architecture fundamentally changes the speculative decoding calculus. The expert activation explosion during verification—where k+1 tokens activate 30-60 unique experts per layer instead of the usual 8—means that speculative decoding for MoE models is not a straightforward extension of dense-model techniques. The MoE-Spec paper's predictions were empirically validated.

Conclusion

The arc from the AllReduce profiling revelation through the n-gram experiment to the EAGLE-3 training pipeline represents a microcosm of the ML infrastructure optimization process at the frontier. It is a story of intellectual honesty—confronting hard hardware constraints, testing hypotheses even when theory suggests they will fail, and methodically building toward the most promising solution despite integration challenges.

The n-gram experiment failed, but that failure was itself a success: it validated the research, confirmed the bottleneck analysis, and cleared the path for the more ambitious EAGLE-3 training pipeline. The pipeline itself, while blocked by API incompatibilities at the segment's end, represented a complete and documented plan for training a custom speculative decoding head—the same approach that Baseten used to achieve 340+ tok/s on Kimi-K2.5.

In the broader narrative of this coding session, this arc demonstrates the value of systematic methodology: profile to find the bottleneck, research the solution space, test the cheapest hypothesis first, document the ambitious path in parallel, and iterate based on empirical results. The path from 82.5 tok/s to 340+ tok/s is not a straight line, but with methodology like this, it is a navigable one.

References