Unblocking the EAGLE-3 Pipeline: How Hidden State Extraction Was Resurrected for Kimi-K2.5 on Blackwell GPUs

Introduction

In the high-stakes world of training speculative decoding models for trillion-parameter language models, the difference between a working pipeline and a blocked one often comes down to a cascade of tiny, individually trivial fixes — each one discovered through hours of methodical source code reading, hypothesis testing, and distributed system debugging. This article chronicles the work of Segment 21 in an ongoing opencode session: the effort to unblock hidden state extraction for the EAGLE-3 training pipeline on the Kimi-K2.5 INT4 model, a 1-trillion-parameter Mixture-of-Experts model deployed across eight NVIDIA RTX PRO 6000 Blackwell GPUs.

The central challenge was an API incompatibility chasm between two rapidly evolving open-source projects. The speculators v0.3.0 library, which provides the hidden state generation infrastructure for EAGLE-3 training, was written against an earlier version of vLLM. The installed environment used vLLM 0.16 nightly — a version that had introduced sweeping architectural changes including a new two-phase execution model, redesigned KV cache management, modified constructor signatures for core classes like Scheduler and Request, and subtle changes to distributed communication primitives like collective_rpc. Bridging this gap required the assistant to methodically discover, understand, and patch each incompatibility, often by reading the actual vLLM source code to reverse-engineer the new API contracts.

The journey spans dozens of messages, multiple extraction attempts, and fixes that range from single-line API patches to complete rewrites of custom model instrumentation code. The debugging arc follows a familiar pattern in distributed ML engineering: each fix unblocks the pipeline to reveal the next, deeper bug, until finally the output is correct — only for the assistant to discover that "correct output" was itself a mirage, masking a subtle distributed communication protocol mismatch that silently corrupted the data.

The State of Play: A Pipeline Blocked by API Evolution

The EAGLE-3 training pipeline consists of several sequential steps. Step 2 — hidden state extraction — is the critical data-generation phase that produces the training targets for the draft model. Without correctly extracted hidden states, the entire speculative decoding training pipeline is blocked. At the start of Segment 21, the pipeline was stuck at exactly this point.

The speculators library's hidden state extraction infrastructure works by monkey-patching the vLLM model's forward pass to intercept intermediate layer activations. It uses a "custom worker" — a specialized vLLM worker class that overrides the model loading and forward execution to capture hidden states at specified layer indices. This custom worker communicates with a generator class that orchestrates the extraction across multiple GPUs using vLLM's distributed execution framework.

The problem was that speculators v0.3.0 had been written against vLLM's API surface as it existed several months earlier. In the intervening period, vLLM had undergone significant architectural evolution. The assistant had already fixed several surface-level issues before Segment 21 began: adding trust_remote_code=True to the tokenizer initialization, adding is_encoder_decoder=False to the SchedulerConfig constructor, and patching the custom worker to navigate the Kimi-K2.5 multimodal wrapper hierarchy (KimiK25ForConditionalGeneration → language_model → model).

But these fixes only scratched the surface. The deeper incompatibilities — changes to function signatures, constructor parameters, execution flow, and distributed communication semantics — remained to be discovered and resolved.

The Cascade of API Incompatibilities

The KV Cache Config Mismatch

The first major API incompatibility discovered in this segment was in the get_kv_cache_config_from_groups() function. In vLLM 0.16, this function had dropped the kv_cache_specs keyword argument that the speculators code was passing. The assistant confirmed this by reading both the speculators call site and vLLM's own internal call to the function, then applied a surgical patch to remove the obsolete parameter.

Rather than immediately re-running the extraction — which would have cost 18+ minutes of model loading time — the assistant made a strategic decision that would prove prescient: it proactively audited the API surface of other vLLM classes that the speculators library used directly. Using Python's inspect.signature() introspection tool, the assistant discovered two additional mismatches: the Request constructor no longer accepted eos_token_id, and the Scheduler constructor now required a block_size parameter that the speculators code did not provide. These two fixes were applied together in a single efficient patch operation, saving what would have been another 18-minute model loading cycle if discovered through trial-and-error.

This "fix one, scan for more" pattern became a hallmark of the assistant's debugging methodology throughout this segment. Each discovered incompatibility triggered a proactive audit of related APIs, compressing what could have been hours of iterative debugging into a single, efficient patch cycle.

The DeepseekV2 Architecture-Specific Bugs

The next layer of bugs was fundamentally different from the API mismatches. These were architectural incompatibilities between the speculators library's generic model instrumentation code and the specific design of the DeepseekV2 model that underlies Kimi-K2.5. The extraction pipeline uses a "custom worker" — a monkey-patched forward pass that intercepts intermediate layer activations — and this worker had been written with assumptions that did not hold for DeepseekV2.

The first clue came from empty error messages. The extraction script was catching exceptions and printing them with print(f"ERROR: {e}"), but e evaluated to an empty string — a notoriously difficult debugging scenario where the error object exists but its string representation is blank. After adding full traceback logging, the assistant paused before re-running the 18-minute model load and decided to investigate the likely cause proactively.

The investigation led to the vLLM model implementation files. The assistant discovered that the DeepseekV2 decoder layer forward signature is:

forward(self, positions, hidden_states, residual, llama_4_scaling=None)

This is a positional-arguments-in-a-specific-order signature with a residual parameter and a llama_4_scaling parameter — neither of which the generic custom worker was handling. The existing patched forward was calling layer(hidden_states=hidden_states, positions=positions, residual=residual) using keyword arguments, but the layer's forward method is defined with positional parameters. When called with keyword arguments, Python maps them to the positional parameters in order — so hidden_states=hidden_states would be passed as the first positional argument, which maps to positions in the layer's signature. This was a catastrophic mismatch that would silently corrupt the data flowing through the model.

Even more critically, the assistant discovered that the DeepseekV2 model uses embed_input_ids() for embedding lookup, not get_input_embeddings() — the method that the speculators library's generic code was calling. This AttributeError was being swallowed by vLLM's distributed error handling, producing the empty error messages that had stymied earlier debugging.

The assistant's response was decisive: rather than attempting incremental patches, it rewrote the entire custom worker from scratch to handle the DeepseekV2 calling convention correctly. The new worker used positional arguments in the correct order, computed llama_4_scaling from the model configuration, and used embed_input_ids() for embedding lookup. It also used the presence of the scoring_func attribute in the model config to detect DeepseekV2 variants and apply the correct calling convention.

The Block Hasher and the Two-Phase Execution Model

With the custom worker fixed, the next extraction attempt loaded the model successfully — only to crash with an AssertionError in the scheduler's block pool: assert len(request.block_hashes) >= num_full_blocks. This was a vLLM 0.16-specific requirement: the KV cache block pool unconditionally expects block hashes on every Request object, even when prefix caching is disabled. The speculators library was creating Request objects without a block_hasher function, leaving block_hashes as an empty list.

The assistant traced this requirement through vLLM's source code, reading the Request constructor, the EngineCore initialization, the block pool assertion, and the scheduler's scheduling logic. The investigation revealed a hidden dependency: the get_request_block_hasher utility function depends on a global variable NONE_HASH that must be initialized by calling init_none_hash() before use. This is a classic example of a global initialization ritual that is easy to overlook when using internal APIs outside their normal context. The assistant updated the patch to include this initialization call, and the block hasher fix was complete.

But the next extraction attempt revealed yet another architectural change. vLLM 0.16 had introduced a two-phase execution model: execute_model() always returns None and stores internal state in self.execute_model_state, and a separate sample_tokens() method must be called to produce the actual output. This was not an optional feature — it was a fundamental redesign of the execution flow. The speculators library, written for the older single-phase model, only called execute_model() and never called sample_tokens(), causing a RuntimeError on the second iteration: "sample_tokens() must be called after execute_model() returns None."

The assistant's initial fix was to disable async_scheduling=False in the SchedulerConfig, hoping this would restore single-phase behavior. But further investigation revealed that the two-phase flow is unconditional in vLLM 0.16 — execute_model() always returns None regardless of the async scheduling flag. The correct fix was to add a sample_tokens(None) call after every execute_model() call, satisfying the state machine protocol even though the sampling output itself was not needed for hidden state extraction.

The Breakthrough and the Mirage

With all fixes applied — the KV cache config patch, the Scheduler and Request constructor patches, the custom worker rewrite, the block hasher initialization, and the two-phase execution fix — the fifth extraction attempt ran to completion. The assistant's celebration was genuine and justified: "IT WORKED! All 10 samples extracted successfully! Hidden states extraction is fully working." The log showed processing at approximately 1,931 tok/s, with the actual extraction taking only 2 seconds after the model finished loading.

But the celebration was premature. When the assistant inspected the actual tensor contents, the data was catastrophically wrong. Instead of 4 tensors (one per captured layer at indices 2, 30, 58, and 60), each of shape [seq_len, 7168] — the full hidden dimension of the model — the saved file contained 771 tensors, each of shape [512]. The hidden dimension had completely vanished. The data was essentially useless for training.

The assistant's response to this discovery exemplifies disciplined debugging. Rather than panicking or making random changes, it systematically analyzed the numbers. The count 771 matched the cumulative token count of the first two samples in the batch (512 + 259 tokens). The shape [512] matched the sequence length of the first sample. This told the assistant that the data was being flattened across the batch dimension and the hidden dimension was being lost entirely.

The assistant traced through the entire data flow — from the _patched_forward that captures (hidden_states + residual) at each target layer, through _store_captured_states that accumulates per-layer lists, through _get_captured_states that concatenates along dimension 0, through the generator's generate() method that slices per sample — to identify where the shape transformation went wrong.

The [0] Indexing Bug: A Single Character That Nearly Broke the Pipeline

The root cause was a subtle API mismatch in vLLM's collective_rpc mechanism. The speculators code contained this pattern:

captured_states_list = self.executor.collective_rpc(
    "_get_captured_states",
    unique_reply_rank=0,
)
aux_hidden_states = captured_states_list[0]

The developer had assumed that collective_rpc always returns a list of results from all worker ranks, and that indexing with [0] extracts the result from rank 0. But as the assistant discovered by reading the vLLM source code via a targeted sed command, the collective_rpc method's docstring states: "Returns single result if unique_reply_rank and/or kv_output_aggregator is provided, otherwise list."

When unique_reply_rank=0 is set, vLLM bypasses the list wrapping and returns the raw result from rank 0 directly. So captured_states_list was already the return value of _get_captured_states — a Python list of four tensors, one per captured layer. The [0] indexing then took only the first layer's tensor (shape [771, 7168]), discarding the other three layers. The subsequent iteration over this single tensor produced 771 slices of shape [512] — exactly the garbage output that was observed.

The fix was a single-character change: removing the [0] indexing. But finding it required tracing data through the entire distributed pipeline, instrumenting both the worker and generator sides with print() statements (since Python's logging module suppressed INFO messages in worker processes), and reading the framework's source code to confirm the API contract.

This bug is particularly instructive because it represents a class of errors that are invisible to any single component. The worker correctly produced four tensors of shape [771, 7168]. The generator correctly received a list of four tensors. But the [0] indexing — which would have been correct if collective_rpc returned a list of per-rank results — silently discarded three-quarters of the data. Neither component was wrong in isolation; they simply disagreed about the contract between them.

The Verification: Extraction Succeeds

With the fix deployed, the assistant launched a re-run of hidden state extraction on 10 test samples. The results were exactly what was needed:

The Cleanup: Consolidating Patches and Removing Instrumentation

With the pipeline unblocked, the assistant did something that distinguishes disciplined engineering from mere problem-solving: it cleaned up. The assistant created and executed clean_debug_prints.py, a script that surgically removed all debug instrumentation from the patched files.

The debug prints — inserted as raw print() calls because worker processes suppressed logger.info() messages — were invasive instrumentation that would clutter production logs and potentially introduce performance overhead. The cleanup script removed them from both custom_worker.py and vllm_hidden_states_generator.py, leaving the substantive bug fixes intact while stripping away the debugging scaffolding.

The assistant's decision to write a dedicated Python script rather than using ad-hoc sed commands reflects a methodical approach. A script could perform context-aware edits, removing only the debug-specific prints while preserving the critical patches. Running it remotely via SSH ensured the cleanup happened on the target machine where the files actually lived.

The Transition: From Extraction to Training

Before the training script could be tested, the environment needed preparation. The assistant executed a compound command that performed four operations in sequence: kill all Python processes, wait for cleanup, release GPU device file handles, and verify memory. The output showed all eight GPUs with 0 MiB used — a clean slate.

This moment marks the boundary between two fundamentally different phases: the conclusion of a debugging campaign focused on data generation, and the beginning of a new phase focused on model training. The eight lines of 0, 0 MiB are not just a status report — they are a blank canvas.

With the GPUs free, the assistant turned to the training script itself. Reading /root/eagle3-train/04_train.py via SSH revealed its purpose:

Step 4: Train EAGLE-3 draft model using speculators. Uses the hidden states extracted in step 2 to train a single-layer Llama-style EAGLE-3 draft model. The training uses TTT (Test-Time Training) where the draft model learns to autoregressively predict the next token from the target model's hidden states. This script does NOT need the target model loaded — only the hidden states files and the verifier's embedding + lm_head weights (extracted separately).

This reading is not casual — it is a deliberate reconnaissance operation. The assistant needs to understand what the script expects, what arguments it takes, and what dependencies it has, before committing to a potentially costly training run that could fail after minutes of execution.

Verifying the Foundation: Import Checks and API Signatures

The assistant then performed a systematic verification of the training infrastructure, testing each import dependency:

  1. speculators.models.eagle3.config — Contains Eagle3SpeculatorConfig, the configuration class needed for model construction. Import succeeds.
  2. speculators.models.eagle3.core.Eagle3DraftModel — The draft model class itself. Import succeeds.
  3. speculators.train.data and submodules — Data loading utilities. Import succeeds.
  4. Specific data classes — Eagle3SampleFileDataset, create_collate_fn, standardize_data_v1, split_files. All import successfully.
  5. Noise transforms — AddUniformNoise, TransformTensors. Also import successfully. But the assistant went deeper. Using Python's inspect.signature() to examine the actual API of Eagle3DraftModel:
__init__ params: ['self', 'config', 't2d', 'd2t']
forward params: ['self', 'hidden_states', 'input_ids', 'lengths', 'loss_mask', 
                 'position_ids', 'verifier_last_hidden_states', 'ttt_steps', 
                 'ttt_step_loss_decay', 'use_off_policy_tokens', 'kwargs']

This introspection revealed that the constructor requires an Eagle3SpeculatorConfig object plus t2d and d2t tensors (token-to-draft and draft-to-token mapping tables). The forward pass expects a rich set of inputs including hidden states, token IDs, sequence lengths, loss masks, position IDs, verifier hidden states, and TTT-specific parameters. This information is critical — the training script must provide all of these correctly, or it will fail at runtime.

Discovering the Built-in Trainer

The assistant also discovered that the speculators library provides a full training infrastructure with TrainerConfig, a train_epoch method, checkpointing, distributed batch sampling, and logging. This discovery raised a strategic question: should the assistant fix the existing custom training script, or pivot to using the library's built-in Trainer?

The assistant searched for CLI entry points to determine if training could be invoked via a simple shell command. The find command returned only __main__.py — no dedicated training CLI existed. This negative result was itself valuable knowledge: the assistant would need to engage with the library's Python-level API, not a turnkey command.

The Silence That Speaks

One of the most structurally interesting moments in this segment is an empty message — a pair of XML-like data tags with nothing between them. Yet this empty message is the system's heartbeat in the autonomous coding session. It is the mechanism by which the opencode runtime signals to the assistant: "Your previous round of tool calls has completed. All results are available. You may now produce your next response."

The assistant's response to this empty signal was one of the longest and most structured messages in the session — a comprehensive status document cataloging hardware specifications, software versions, model architecture details, profiling results, the complete list of patches applied, the status of every pipeline step, and a prioritized list of next steps. The empty message triggered this synthesis because the assistant recognized that it had reached a natural transition point: the debugging phase was complete, hidden state extraction was verified working, and the next major phase required a different approach.

Lessons for Debugging Distributed ML Systems

The debugging journey chronicled in this segment offers several enduring lessons for engineers working with complex ML infrastructure:

1. Read the source code. The single most powerful debugging technique demonstrated throughout this session was reading the actual source code of the libraries involved. The assistant repeatedly used grep, sed, and cat to inspect vLLM's internal implementation files, discovering function signatures, constructor parameters, and execution patterns that were not documented anywhere else. When error messages are empty or misleading, the source code is the ground truth.

2. Proactive auditing beats reactive debugging. Rather than fixing one error and re-running the expensive extraction (which took 18+ minutes for model loading), the assistant proactively audited the API surface of related classes after finding each bug. This "fix one, scan for more" pattern saved hours of wasted computation.

3. Distributed systems have hidden initialization rituals. The init_none_hash discovery is a perfect example of a global initialization dependency that is invisible when using APIs outside their normal context. When integrating with internal framework APIs, always check for module-level state that must be initialized before use.

4. "It ran without errors" is not the same as "it produced correct output." The most dangerous bugs are those that produce plausible-looking but wrong results. The assistant's discipline in verifying tensor shapes — even after celebrating a successful run — caught a subtle distributed communication bug that would have silently corrupted the entire training dataset.

5. Architecture-specific code cannot be generic. The DeepseekV2 model's unique forward signature, with its positional arguments, residual stream, and llama_4_scaling parameter, required a completely custom instrumentation approach. Generic hooks designed for Llama-style architectures simply cannot handle the diversity of modern LLM designs.

6. Distributed debugging requires multi-sided instrumentation. The [0] bug was invisible from either side alone. The worker reported correct captures; the generator reported wrong output. Only by instrumenting both ends simultaneously could the assistant identify the mismatch.

7. Cleanup is part of debugging. The assistant didn't celebrate the extraction breakthrough and move on. It cleaned debug instrumentation, consolidated patches, freed resources, and documented progress. This transforms a fragile sequence of hacks into a maintainable pipeline.

8. The transition between phases is itself a phase. The GPU cleanup command, the reading of the training script, the import checks, the API introspection — these are not distractions from the "real work." They are the work. A clean handoff between pipeline stages, with all state verified and all assumptions checked, is the difference between a reproducible system and a brittle sequence of lucky accidents.

Conclusion

Segment 21 of this opencode session tells a complete story: from the breakthrough of fixing the cascade of API incompatibilities that had blocked hidden state extraction, through the discovery and resolution of the insidious [0] indexing bug that silently corrupted the output, through the verification that extraction was producing correct tensors, the disciplined cleanup of debug instrumentation, the methodical transition to the training phase, and the reconnaissance of the training infrastructure.

The hidden state extraction pipeline is now operational. The training infrastructure has been mapped. The imports are verified. The GPUs are clean. The next step — actually training the EAGLE-3 draft model — awaits. But before that can happen, the assistant must understand the built-in Trainer's API, ensure data format compatibility, and write the orchestration code. That work begins in the next phase of the session.

The debugging journey chronicled here is a testament to the complexity of integrating bleeding-edge ML frameworks. The assistant faced not one but five distinct classes of bugs: straightforward API mismatches (KV cache config, Scheduler, Request constructors), architecture-specific code incompatibilities (DeepseekV2 forward signature, embed_input_ids), distributed system initialization rituals (the init_none_hash global), fundamental execution model redesigns (the two-phase execute_model/sample_tokens flow), and subtle distributed communication protocol mismatches (the collective_rpc unique_reply_rank indexing error).

Each fix was individually small — a removed keyword argument, an added parameter, a rewritten function, a corrected indexing operation. But discovering each one required reading source code, understanding the architectural assumptions of both the speculators library and vLLM 0.16, and methodically tracing through distributed execution paths. The assistant's willingness to fork and modify code without hesitation, its systematic approach to hypothesis generation and testing, and its discipline in verifying outputs even when the pipeline ran without errors — these are the skills that ultimately unblocked the EAGLE-3 training pipeline and made speculative decoding for a trillion-parameter model a step closer to reality.