The Hidden State Mismatch: A Debugging Breakthrough in EAGLE-3 Speculative Decoding

Introduction

In the complex world of speculative decoding for large language models, few things are more frustrating than a model that trains perfectly but performs poorly at inference time. This was precisely the situation facing the developer of an EAGLE-3 draft model for the Kimi-K2.5 architecture. The draft model had achieved a respectable 74.7% validation accuracy during training, yet when deployed with SGLang for speculative decoding, it delivered a paltry ~56.8 tok/s against a 90.0 tok/s baseline — barely better than no speculation at all. Message 4458 in this coding session marks the critical turning point: the moment when the developer traced the root cause to a fundamental mismatch between how hidden states were constructed during training versus how they were being fed to the model during inference.

The Context of the Investigation

The developer had been working for days on deploying an EAGLE-3 draft model for the Kimi-K2.5 large language model. The pipeline involved training a lightweight "draft" model that predicts multiple future tokens in parallel, which the target model then verifies. When the draft model's predictions are accurate, speculative decoding can dramatically speed up inference. However, the developer's draft model was consistently underperforming, with an accept length of only ~1.6 out of 16 possible draft tokens — meaning the target model was rejecting almost all of the draft's predictions.

Earlier in the session (see [msg 4441]-[msg 4454]), the developer had written a standalone test to isolate the draft model from the SGLang inference engine, hoping to determine whether the issue was in the model weights or in the inference wiring. That test revealed a shocking discrepancy: the draft model achieved only 34.1% accuracy on training data when run standalone, compared to the 74.7% reported during training. This massive gap pointed to a fundamental difference between the training and inference pipelines.

The Critical Discovery

Message 4458 begins with the developer reading the training data pipeline code in the speculators library. The key function is shift_batch in speculators/train/data.py, which prepares batches for the EAGLE-3 training loop. The developer examines lines 125-128, which handle the eagle3 format specifically:

"hidden_states": torch.cat(data["hidden_states"][:-1], dim=-1),  # all but last
"verifier_last_hidden_states": data["hidden_states"][-1],         # last one

This is the moment of revelation. The data["hidden_states"] is a list of four tensors: [embed_output, layer_3, layer_31, layer_59]. These represent the hidden states captured at specific layers of the target model during the forward pass. The training code concatenates the first three of these — cat([embed, layer3, layer31], dim=-1) — and uses that as the input to the draft model's fully-connected (fc) projection layer. The last one (layer59) is stored separately as verifier_last_hidden_states for the verifier model.

But the SGLang inference pipeline was doing something completely different. It was concatenating cat([layer3, layer31, layer59]) — the last three hidden states, excluding the embedding output entirely. This meant the draft model was receiving a completely different input representation than what it was trained on.

Why This Matters

The EAGLE-3 architecture uses a technique called "feature-level speculative decoding" where the draft model doesn't just predict tokens from the target model's final hidden states — it also uses intermediate representations from multiple layers. The idea is that these intermediate representations contain rich information about the target model's processing that can help the draft model make more accurate predictions.

The training pipeline captured four hidden states:

  1. Embedding output (layer index -1): The raw token embeddings before any transformer processing
  2. Layer 3: Early transformer layer representations
  3. Layer 31: Mid-to-late transformer layer representations
  4. Layer 59: Near-final transformer layer representations During training, the model learned to use the first three of these (embed, layer3, layer31) to predict the next tokens. The embedding output is particularly important because it provides the raw token identity information — knowing which token is at each position is essential for predicting what comes next. By excluding the embedding output and instead including layer59 (which the model had never seen during training), the inference pipeline was feeding the draft model a completely alien input representation.

The Reasoning Process

The developer's thinking in this message is a textbook example of systematic debugging. The chain of reasoning goes:

  1. Hypothesis: The training and inference pipelines might construct hidden states differently.
  2. Evidence: The standalone test showed 34.1% accuracy vs 74.7% training accuracy, confirming something is different.
  3. Investigation: Read the training data code to understand exactly how hidden states are prepared.
  4. Discovery: Lines 125-128 reveal the [:-1] slice — training uses all but the last hidden state.
  5. Verification: Cross-reference with the inference code to confirm the mismatch. The developer also notices a secondary detail: the training data is shifted. input_ids[1:] (removing the first token) becomes the prediction targets, while hidden_states[:-1] (removing the last position) provides the inputs. This is the standard autoregressive shift — at position i, the model uses hidden states from position i to predict the token at position i+1. The developer notes this but recognizes it's a separate concern from the concatenation order issue.

Assumptions and Mistakes

Several assumptions had led to this bug:

  1. The assumption that hidden states were ordered by layer index: The developer had assumed that data["hidden_states"] was ordered [layer3, layer31, layer59] (increasing layer index), but it was actually [embed, layer3, layer31, layer59] with the embedding output prepended.
  2. The assumption that all captured hidden states should be used: The training code explicitly excludes the last hidden state (layer59) from the draft model's input, using it only for the verifier. The inference code was including all three "auxiliary" hidden states without realizing one was reserved.
  3. The assumption that the embedding output was not needed: The developer had configured the auxiliary hidden state capture with layer IDs [2, 30, 58] (0-indexed, corresponding to layers 3, 31, 59 in 1-indexed notation), not realizing that the embedding output (layer -1) was also being captured and prepended to the list.

Input Knowledge Required

To fully understand this message, the reader needs:

  1. EAGLE-3 architecture knowledge: Understanding that EAGLE-3 uses multiple hidden states from the target model as input to the draft model, concatenated along the feature dimension.
  2. Speculative decoding concepts: How draft models propose tokens and target models verify them, and why the input representation matters.
  3. The Kimi-K2.5 model structure: That it has 60+ transformer layers, and hidden states are captured at specific layers for speculative decoding.
  4. Python tensor operations: Understanding cat([...], dim=-1) (concatenation along the last dimension) and [:-1] slicing.
  5. The training data pipeline: How shift_batch prepares training examples by shifting input_ids and hidden_states for autoregressive prediction.

Output Knowledge Created

This message creates several critical pieces of knowledge:

  1. The root cause of the performance gap: The hidden state concatenation order mismatch explains why the draft model was performing poorly despite good training metrics.
  2. The correct hidden state format: The draft model expects cat([embed_output, layer_3, layer_31]) — three tensors of shape [seq_len, 7168] concatenated to form [seq_len, 21504].
  3. The verifier's hidden state: Layer 59 is reserved for the verifier model, not the draft model.
  4. The data shift pattern: Training uses input_ids[1:] as targets and hidden_states[:-1] as inputs, matching the standard autoregressive formulation.

Impact and Resolution

This discovery directly led to the fix implemented later in the session: modifying the SGLang inference code to capture the embedding output (using layer_id=-1) and changing the draft model configuration from [2, 30, 58] to [-1, 2, 30] to include the embedding in the correct position. After the fix, the standalone test accuracy jumped to 76.9%, matching the training metrics. However, the overall inference speed only improved to 54.8 tok/s — still far below the 90.0 baseline — suggesting that while this was a critical bug, other issues remained in the inference pipeline.

Conclusion

Message 4458 represents the classic debugging breakthrough: the moment when scattered symptoms suddenly cohere into a single, testable hypothesis. The developer's systematic approach — isolating the problem with a standalone test, tracing the data pipeline in both training and inference code, and comparing them line by line — exemplifies the methodical thinking required to debug complex ML systems. The hidden state mismatch is a cautionary tale about the subtle ways that training and inference pipelines can diverge, and a reminder that even well-trained models will fail if they receive the wrong inputs at deployment time.