Reading the Blueprint: How the Assistant Uncovered vLLM's Hidden State Extraction Pipeline for EAGLE-3 Training

Introduction

In the long arc of optimizing Kimi-K2.5 inference on 8× Blackwell GPUs, a pivotal moment arrived when the assistant, having ruled out n-gram speculative decoding as a viable optimization, turned to building a custom EAGLE-3 training pipeline. The message at <msg id=2520> captures a seemingly mundane action: reading a single Python file from the installed speculators library on the remote machine. But this act of reading was the culmination of hours of research, experimentation, and strategic decision-making. It represents the moment when the assistant shifted from theoretical exploration to concrete implementation, choosing to understand the exact mechanism of hidden state extraction rather than blindly adapting an incompatible library.

The Context: Why This File Matters

To understand the significance of this message, one must appreciate the problem the assistant was solving. The profiling session had revealed that AllReduce communication dominated 51.5% of decode time on the 8-GPU setup. Speculative decoding — generating multiple draft tokens cheaply and verifying them in parallel — promised to improve throughput by hiding the AllReduce latency behind additional computation. But the assistant had already proven that n-gram speculation was counterproductive: it made throughput worse by 9–26% because the MoE verification overhead of loading extra experts for rejected tokens outweighed the rare n-gram matches in a reasoning model's novel thinking chains.

The only viable path forward was training a custom EAGLE-3 draft model specifically for Kimi-K2.5. The user had directed the assistant to "start implementing the training scripts, on our existing machine with lowered numbers, then should be easy to port to much more expensive b300 machine for a simple train run." This directive set the stage for the assistant's work in <msg id=2520>.

The assistant had already completed extensive research in parallel task agents: exploring the SpecForge framework (SGLang's training library), the Speculators framework (vLLM's training library), and the existing AQ-MedAI/Kimi-K2-Instruct-eagle3 draft model on HuggingFace. The decision had been made to use Speculators because it was part of the vLLM ecosystem that the deployment already used. Speculators had been installed successfully (speculators==0.3.0), and the core imports (Eagle3DraftModel, VllmHiddenStatesGenerator, HiddenStatesWorkerExtension) all worked. But there was a critical unknown: Speculators was designed for vLLM ≤0.15, while the running deployment used vLLM 0.16.0rc2. The internal API surface might have changed.

What the Message Actually Shows

The message is a single bash command executed over SSH:

ssh root@10.1.230.174 'cat ~/ml-env/lib/python3.12/site-packages/speculators/data_generation/vllm_hidden_states_generator.py'

The assistant then displays the file contents, which begin with:

"""Extract hidden states from intermediate layers during prefill using vLLM."""

import torch
from transformers import AutoConfig, AutoTokenizer
from vllm.config import (
    CacheConfig,
    DeviceConfig,
    LoadConfig,
    ModelConfig,
    ParallelConfig,
    SchedulerConfig,
    VllmConfig,
)
from vllm.sampling_params import SamplingParams
from vllm.v1.core.kv_cache_utils import (
    _get_kv_cache_groups_uniform_spec,
    get_kv_cache_config_from_groups,
    unify_hybrid_kv_cache_specs,
)

The file is 306 lines long. The assistant is reading the entire file to understand the exact API surface that Speculators uses to hook into vLLM's forward pass and extract intermediate hidden states.

The Reasoning Behind Reading This File

The assistant's thinking process, visible in the preceding messages, reveals a careful cost-benefit analysis. There were three approaches to generating training data (hidden states from Kimi-K2.5's intermediate layers):

Approach A: Use Speculators' VllmHiddenStatesGenerator directly. This was the cleanest path but risked runtime errors because the library monkey-patches vLLM internals that may have changed between v0.15 and v0.16.

Approach B: Write a standalone script using HuggingFace Transformers with register_forward_hook to capture hidden states. This would be framework-independent but would require loading the 1T-parameter model without INT4 Marlin kernel acceleration, making it impractically slow.

Approach C: Write a custom script that uses vLLM's offline LLM class with the Speculators worker extension. This was the middle ground — using vLLM's optimized inference but adapting the extension code to work with v0.16.

Before committing to any approach, the assistant needed to understand exactly what VllmHiddenStatesGenerator does internally. The file being read reveals the key design: it imports vLLM configuration classes (CacheConfig, DeviceConfig, ModelConfig, ParallelConfig, SchedulerConfig, VllmConfig), sampling parameters, and KV cache utilities. These are the internal APIs that may have changed between versions.

Assumptions Made

The assistant made several assumptions in this message:

  1. That reading the source code directly would reveal API compatibility issues faster than trial-and-error. This was a reasonable assumption — understanding the exact imports and function signatures would let the assistant predict whether v0.16 changes would break the pipeline.
  2. That the Speculators library's data generation path is the right one to adapt. The assistant had already committed to Speculators over SpecForge, and reading this file was part of understanding how to make it work with v0.16 rather than abandoning it.
  3. That the file contains the complete logic for hidden state extraction. The assistant assumed that vllm_hidden_states_generator.py is the primary entry point for data generation, not just a wrapper around deeper logic.
  4. That the running vLLM server could be stopped and restarted for data generation. The assistant had earlier considered adding a custom endpoint to the running server but ultimately decided that offline data generation (stopping the service temporarily) was acceptable.

Input Knowledge Required

To understand this message fully, one needs:

  1. The EAGLE-3 architecture: Understanding that EAGLE-3 training requires extracting hidden states from specific intermediate layers of the target model (in this case, layers 2, 30, and 58 of Kimi-K2.5), then training a lightweight draft model to predict those hidden states autoregressively.
  2. vLLM's internal architecture: Knowledge of classes like ModelConfig, ParallelConfig, SchedulerConfig, and how they configure the model execution pipeline. The VllmConfig class (introduced in v0.16) consolidates these into a single configuration object.
  3. The Speculators framework: Understanding that Speculators provides a complete pipeline: hidden state extraction via vLLM monkey-patching, vocabulary mapping between target and draft models, and training using a custom Eagle3DraftModel with a SpeculatorsTrainer.
  4. The deployment environment: The model is Kimi-K2.5 (a 1T-parameter MoE model) running on 8× RTX PRO 6000 Blackwell GPUs with vLLM 0.16.0rc2, using INT4 quantization via Marlin kernels.
  5. The version mismatch problem: Speculators v0.3.0 was designed for vLLM ≤0.15, and the assistant needs to determine whether the internal APIs it depends on have changed in v0.16.

Output Knowledge Created

This message produces several forms of knowledge:

  1. The exact API surface of Speculators' data generation: The assistant now knows that VllmHiddenStatesGenerator imports SchedulerConfig (not SchedulerConfig from v1), CacheConfig, DeviceConfig, etc. It also uses KV cache utilities from vllm.v1.core.kv_cache_utils, which may have changed signatures.
  2. A baseline for compatibility assessment: By seeing the exact imports and function signatures, the assistant can now compare them against vLLM 0.16's actual API to identify breakpoints.
  3. Documentation of the approach: The assistant is building knowledge that will be recorded in the next-steps-eagle.md document and the eagle3-train/ directory, creating a reproducible pipeline.
  4. Confidence in the implementation path: Reading this file confirms that the Speculators approach is well-structured and that adapting it for v0.16 is feasible — the core logic is in a single file with clear dependencies.

The Thinking Process

The assistant's reasoning in the messages leading up to <msg id=2520> shows a methodical approach to solving the compatibility problem. After installing Speculators and verifying that imports work at the Python level, the assistant recognized that import-time success doesn't guarantee runtime success. The monkey-patching in custom_worker.py and the configuration initialization in vllm_hidden_states_generator.py could fail if vLLM's internal APIs changed.

The assistant's thought process was: "Let me check if we can use vLLM offline LLM with the worker extension. We don't need the server running simultaneously. We stop the server, run data generation, then restart. Or even better — let me check if we can add a custom endpoint to the running server that returns hidden states. Actually, the cleanest approach: Write a script that uses vLLM's offline LLM class with the speculators worker extension. This is exactly what the speculators data_generation_offline.py does. Let me just read that script and adapt it for our model."

This reasoning chain shows the assistant weighing options, rejecting the more complex approaches (custom server endpoint, standalone Transformers loading), and converging on the simplest adaptation path: using Speculators' existing code but understanding it well enough to fix any v0.16 incompatibilities.

The Broader Significance

This message, while seemingly minor, represents a critical transition in the session. The assistant had spent hours profiling, researching, and experimenting. Now it was time to build. Reading the source code was the first concrete step toward implementation — understanding the exact mechanism before modifying it.

The file revealed something important: Speculators' data generation depends on specific vLLM configuration classes. If v0.16 changed SchedulerConfig (which it did — v0.16 introduced VllmConfig as a wrapper), the data generation would fail. The assistant would need to patch the imports and initialization logic.

This reading session also demonstrated a key principle of the assistant's methodology: never trust documentation when source code is available. Rather than relying on Speculators' README or API docs, the assistant went straight to the implementation to understand the exact dependencies and mechanisms.

Conclusion

Message <msg id=2520> captures the moment when the assistant moved from research to implementation in building an EAGLE-3 training pipeline for Kimi-K2.5. Reading the vllm_hidden_states_generator.py file was not a passive act — it was a deliberate, strategic decision to understand the exact API surface that would need to be adapted for vLLM 0.16 compatibility. The assistant's methodical approach — research both frameworks, install the chosen one, verify imports, read source code, then implement — reflects a disciplined engineering practice that would prove essential when the inevitable API mismatches were discovered in subsequent messages. This single cat command, executed over SSH, was the foundation upon which the entire EAGLE-3 training pipeline would be built.