The Critical Script: Building the EAGLE-3 Data Generation Pipeline for Kimi-K2.5

A Pivot from Profiling to Speculative Decoding

The message at index 2524 in this opencode session is a deceptively simple one-line action: the assistant writes a file called 01_prepare_dataset.py into a newly created eagle3-train/ directory. But this single file write represents the culmination of an extensive multi-threaded research phase, a strategic pivot from bottleneck profiling to throughput optimization, and a series of careful architectural decisions about how to train a custom EAGLE-3 draft model for the 1-trillion-parameter Kimi-K2.5 model running on 8×RTX PRO 6000 Blackwell GPUs.

The full text of the message reads:

[assistant] Now the data generation script. This is the most critical piece — it needs to stop the vLLM server, use speculators' VllmHiddenStatesGenerator to extract hidden states, and save them in the right format: [write] /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/01_prepare_dataset.py Wrote file successfully.

To understand why this message was written, we must trace the reasoning chain that led to it. The session had spent dozens of messages conducting a comprehensive profiling campaign of the Kimi-K2.5 INT4 deployment, culminating in the discovery that AllReduce communication consumed 51.5% of decode time — a bottleneck fundamentally limited by PCIe bandwidth between the 8 GPUs. With hardware-bound optimizations exhausted, the user and assistant pivoted to speculative decoding as a software-only path to higher throughput.

The Research Phase: Why EAGLE-3 Won

Before writing a single line of pipeline code, the assistant launched three parallel research task agents ([msg 2507]) to investigate the landscape. These agents explored the SpecForge framework (SGLang's training library), the Speculators library (vLLM's training library), and the existing AQ-MedAI/Kimi-K2-Instruct-eagle3 draft model on HuggingFace. The research revealed critical insights:

  1. N-gram speculation was a dead end. Empirical testing showed it was 9–26% slower than baseline inference because MoE expert activation overhead during verification of rejected draft tokens outweighed any gains from the rare n-gram matches (17–31% acceptance rate). Reasoning models generate novel thinking chains with little repetition, making n-gram approaches fundamentally unsuitable.
  2. The only off-the-shelf EAGLE-3 draft model (AQ-MedAI/Kimi-K2-Instruct-eagle3) was trained for Kimi-K2, not K2.5. The distribution mismatch would lower acceptance rates, and the draft vocabulary of 32K tokens didn't match K2.5's larger vocabulary.
  3. The Baseten approach — training a custom EAGLE-3 head on the target model's own hidden states — was identified as the most promising path, consistent with recent research showing 1.6–2.2× speedups for MoE models with custom-trained draft models.

The Decision to Use Speculators

A key architectural decision visible in the reasoning leading up to message 2524 is the choice of framework. The assistant had two options: SpecForge (from the SGLang project) and Speculators (from the vLLM project). After reading both codebases, the assistant chose Speculators because "we're already on vLLM and it has a cleaner pipeline" ([msg 2509]).

However, this decision came with a known risk: Speculators v0.3.0 was designed for vLLM ≤0.15, while the installed environment had vLLM 0.16.0rc2. The assistant explicitly acknowledged this tension in [msg 2516]: "The speculators data generation path may still break at runtime with vLLM 0.16 since it monkey-patches internals. Rather than fighting that, I'll write our own approach." Yet by message 2524, the assistant had reversed this position and decided to use Speculators' VllmHiddenStatesGenerator directly. This reversal is the most interesting aspect of the message.

The Reasoning Behind the Reversal

The assistant's thinking evolved through several stages. Initially, it considered three alternative approaches:

  1. PyTorch hooks on the running server: Capture hidden states by registering forward hooks on the live vLLM server. This would avoid downtime but required deep integration with vLLM's internal execution model.
  2. HuggingFace Transformers standalone: Load the model directly using HuggingFace APIs and register_forward_hook. This was rejected because loading a 1T-parameter model in Transformers would be extremely slow and would lose the INT4 Marlin kernel acceleration that vLLM provides.
  3. Custom vLLM endpoint: Add a custom API endpoint to the running server that returns hidden states. This was the most elegant approach but required modifying vLLM source code. Ultimately, the assistant converged on the pragmatic choice: use Speculators' VllmHiddenStatesGenerator, which spawns a fresh vLLM instance via MultiprocExecutor. This approach required stopping the production server, running data generation, then restarting — but it reused existing, tested code rather than requiring custom implementation of the complex hidden state capture mechanism. The message text itself reveals the assistant's understanding of the tradeoffs: "This is the most critical piece — it needs to stop the vLLM server, use speculators' VllmHiddenStatesGenerator to extract hidden states, and save them in the right format." The assistant is explicitly acknowledging that stopping the server is a necessary cost of this approach.

Assumptions and Potential Pitfalls

Several assumptions underpin this message, some of which would later prove problematic:

Assumption 1: The VllmHiddenStatesGenerator would work with vLLM 0.16. The generator creates its own VllmConfig and MultiprocExecutor, importing from vllm.config modules that may have changed between versions. The assistant had verified that imports succeeded (<msg id=2514, 2515>) but had not tested runtime execution.

Assumption 2: The model architecture path would match. Speculators' monkey-patching hooks into model.model.layers to capture hidden states, but Kimi-K2.5 uses a multimodal wrapper architecture where the actual transformer layers are at model.language_model.model.layers. This mismatch would later cause runtime errors ([msg 2541]).

Assumption 3: The data format would be compatible. The assistant was writing 01_prepare_dataset.py to use HuggingFace datasets to load and tokenize prompts, then pass them to the hidden state generator. The format expected by the speculators trainer (data.py) was .pt files containing dictionaries with specific keys like hidden_states, input_ids, attention_mask, and labels. Getting this format exactly right required careful reading of the training data loader.

Assumption 4: Stopping and restarting the server was acceptable. The vLLM server took approximately 30 minutes to load the 547GB model. Stopping it for data generation meant significant downtime. The assistant implicitly assumed that data generation would be a batch operation that could be completed in a single session.

Input Knowledge Required

To understand this message, one needs:

Output Knowledge Created

This message creates:

The Thinking Process

The assistant's reasoning in this message is best understood by examining the sequence of decisions that preceded it. In [msg 2519], the assistant walked through four different approaches in rapid succession, each rejected in turn:

  1. "Write a standalone hidden state extraction script that loads the model directly using HuggingFace transformers" — rejected because it would be slow and lose INT4 acceleration.
  2. "Use our running vLLM server for inference + write a separate lightweight model" — rejected as unnecessarily complex.
  3. "Use vLLM offline LLM with the worker extension" — this was the chosen approach.
  4. "Add a custom endpoint to the running server" — rejected as requiring too much vLLM modification. By message 2524, the assistant had committed to approach #3 and was implementing it. The phrase "this is the most critical piece" reflects the assistant's understanding that data generation is the bottleneck in the entire training pipeline — the research phase had already established that data generation would take 3–7 days on the current hardware versus 12–24 hours on rented B200 NVL8 machines. The message also reveals the assistant's systematic approach to building the pipeline: it was creating scripts in numbered order (01, 02, 03, 04) with a clear dependency chain. The 01_prepare_dataset.py script was the entry point that all subsequent stages would depend on.

Conclusion

Message 2524 is a pivotal commit point in a complex engineering effort. It represents the moment when research and analysis crystallized into implementation — when the assistant stopped evaluating alternatives and started building. The decision to use Speculators' VllmHiddenStatesGenerator despite known version incompatibilities was a calculated risk, prioritizing code reuse and rapid implementation over robustness. This tradeoff would later prove costly when runtime errors emerged ([msg 2541]), but it was the right decision given the constraints: the assistant was building a pipeline to be ported to rented hardware, and getting a working end-to-end flow on the current machine was the primary goal, even if it required patches and workarounds along the way.