The Hidden State Concatenation Bug: A Eureka Moment in Debugging EAGLE-3 Speculative Decoding

In the middle of a protracted debugging session spanning dozens of messages, one particular message — message index 4459 — stands out as the moment of discovery. It is a short, almost telegraphic message, but it crystallizes the root cause of a frustrating performance problem that had consumed hours of investigation. The assistant had been wrestling with a trained EAGLE-3 draft model that achieved only 54.8 tokens per second in speculative decoding, far below the 90 tok/s baseline, despite reporting 74.7% validation accuracy during training. Something was fundamentally wrong with how the draft model was being served. Message 4459 is the instant when the assistant finally understood what that something was.

The Context of the Investigation

To appreciate the significance of message 4459, one must understand the chain of reasoning that led to it. The assistant had been deploying an EAGLE-3 draft model for the Kimi-K2.5 large language model, using SGLang as the inference engine. The draft model is a small "speculator" network that predicts multiple future tokens in parallel, allowing the main model to verify and accept them, thereby accelerating inference. The training pipeline had produced a draft model with 74.7% validation accuracy on the first predicted token — a respectable figure that should have translated into meaningful speedups.

Yet when the assistant launched the SGLang server with speculative decoding enabled, performance was abysmal. The acceptance rate hovered around 1.6–1.9 tokens per verification step, and overall throughput was roughly half the baseline. Something was clearly broken in the inference pipeline, but the cause was not obvious.

The assistant systematically eliminated possibilities. First, it discovered that --speculative-num-steps 1 was silently overriding --speculative-num-draft-tokens 16 due to a constraint in SGLang when topk=1, limiting the draft to just 2 tokens. Fixing that to --speculative-num-steps 15 actually made performance worse (46.7 tok/s), suggesting the draft model itself was not predicting well despite its training accuracy.

To isolate the draft model from SGLang's complexity, the assistant wrote a standalone test (test_drafter_standalone.py) that loaded the trained weights and evaluated them directly on training data. The result was shocking: only 34.1% accuracy on the very data the model was trained on, compared to the 74.7% reported during training. This confirmed that the inference pipeline was feeding the draft model different inputs than it had been trained on.

The Moment of Discovery

Message 4459 is the assistant's synthesis of this discovery. The message begins with the assistant articulating the full picture it has just pieced together:

Training hidden_states = cat([embed, layer3, layer31], dim=-1) = cat(data["hidden_states"][:-1]) → [seq_len, 3*7168] But our HS dump has 4 entries: [embed, layer3, layer31, layer59] Training uses the first 3 (embed + 2 aux layers), NOT the last 3 (3 aux layers) This means the fc layer was trained on cat([embed_output, layer3, layer31]) but we're feeding it cat([layer3, layer31, layer59]). That's the bug!

This is the core insight. The hidden state dump contains four tensors per sample: the embedding output (layer index -1, conceptually), and the outputs of three auxiliary layers (layers 3, 31, and 59 in the 60-layer DeepSeekV3 model). During training, the standardize_data_v1 function in the speculators training pipeline concatenates the first three of these — data["hidden_states"][:-1] — which corresponds to cat([embed, layer3, layer31]). The fourth hidden state (layer 59) is stored separately as verifier_last_hidden_states.

But SGLang's inference pipeline was capturing only the three auxiliary hidden states (the outputs of layers 3, 31, and 59) and concatenating them — cat([layer3, layer31, layer59]) — completely omitting the embedding output. The draft model's fully-connected (fc) layer, which projects the concatenated input down to the model's hidden dimension, had been trained on a very different input distribution. It was receiving three auxiliary hidden states with large norms (layer 31's norm averaged 174, layer 59's norm averaged 111) when it expected the embedding output (norm ~1) plus two auxiliary states. This explained the catastrophic accuracy drop from 74.7% to 34.1%.

The Reasoning Process Visible in the Message

The assistant's thinking is remarkably clear in this message. It does not just state the bug — it demonstrates how it arrived at the conclusion by cross-referencing two sources of information:

The training data pipeline code (examined in message 4458), which showed that standardize_data_v1 uses data["hidden_states"][:-1] — all but the last hidden state.
The hidden state dump structure, which the assistant had verified contains exactly 4 entries: [embed, layer3, layer31, layer59]. The assistant then immediately takes the next logical step: it checks the training script 04_train.py to confirm which dataset format path was actually used during training. This is crucial because the training pipeline supports multiple data formats (v0, v1, etc.), and the assistant needs to verify that the format it analyzed is indeed the one that was executed. The message ends with a grep command searching for relevant keywords in the training script. This dual verification — understanding both the data format specification and confirming it was actually used — is characteristic of rigorous debugging. The assistant is not satisfied with a hypothesis; it must confirm every link in the chain.

Assumptions and Mistakes

Several assumptions had been made that turned out to be incorrect:

The assumption that SGLang's hidden state capture matched training. The most fundamental mistake was assuming that the hidden states captured by SGLang during inference would have the same structure as those used during training. The SGLang model wrapper (deepseek_v2.py) was configured to capture auxiliary hidden states from specific layer IDs [2, 30, 58] (0-indexed, corresponding to layers 3, 31, 59 in 1-indexed notation). But the training pipeline expected the embedding output to be included as the first element of the hidden states list.

The assumption that the draft model config specified the correct layers. The draft model's eagle_aux_hidden_state_layer_ids was set to [2, 30, 58], which told SGLang which layers to capture. But the training pipeline's standardize_data_v1 function implicitly assumed that the first element of the hidden states list was the embedding output, not an auxiliary layer. The config should have been [-1, 2, 30] to include the embedding output (layer -1) as the first captured state.

The assumption that 74.7% training accuracy would translate to inference. This assumption was reasonable on its face — a model that achieves 74.7% accuracy on held-out validation data should perform similarly when served. But the validation accuracy was measured on correctly formatted inputs (embed + 2 aux layers), while inference was feeding incorrectly formatted inputs (3 aux layers). The accuracy discrepancy (34.1% vs 74.7%) was entirely explained by this input mismatch.

Input Knowledge Required

To fully understand this message, one needs knowledge of several domains:

EAGLE-3 architecture. EAGLE-3 is a speculative decoding framework where a small draft model predicts future tokens using hidden states from the target (verifier) model. The draft model takes as input a concatenation of hidden states from selected layers of the target model, plus the embedding output, and uses a fully-connected layer to project this down to its hidden dimension.

The speculators training pipeline. The training code in the speculators package uses a standardize_data_v1 function that expects hidden states as a list of tensors. It concatenates all but the last element along the last dimension to form the draft model's input, and uses the last element as the verifier's last hidden state.

SGLang's model wrappers. SGLang uses model-specific Python files (like deepseek_v2.py and kimi_k25.py) that define how hidden states are captured and passed to the speculative decoding system. The capture_aux_hidden_states mechanism in SGLang's logits_processor.py collects hidden states from specified layer IDs and concatenates them.

Layer indexing conventions. The DeepSeekV3 model has 60 layers indexed 0–59. The embedding output is conceptually layer -1. The draft model config used 0-indexed layer IDs [2, 30, 58], which correspond to layers 3, 31, and 59 in 1-indexed notation.

Output Knowledge Created

This message creates several important pieces of knowledge:

The precise bug location. The mismatch between training and inference hidden state concatenation is definitively identified. The training pipeline uses cat([embed, layer3, layer31]) while SGLang was using cat([layer3, layer31, layer59]).

The fix direction. The assistant immediately recognizes that the fix requires either (a) modifying SGLang's model wrapper to capture the embedding output as an additional hidden state, or (b) modifying the draft model's config to include layer -1 (embedding) as the first captured state. The subsequent messages show the assistant implementing option (a) by modifying deepseek_v2.py to capture the embedding output when layer_id=-1 is specified, and updating the config from [2, 30, 58] to [-1, 2, 30].

A methodology for debugging speculative decoding. The standalone test approach — bypassing SGLang entirely and evaluating the draft model directly on training data — proves to be an effective debugging technique. It isolates the draft model's intrinsic quality from the inference pipeline's wiring, making it possible to attribute accuracy drops to the correct component.

The Broader Significance

Message 4459 is a textbook example of a common class of bugs in ML systems: training-inference skew. The training pipeline and the inference pipeline must process data identically, but when they are implemented by different codebases (the speculators training framework vs. SGLang's inference engine), subtle differences can creep in. The embedding output — a tensor with norm ~1 — is qualitatively different from the auxiliary hidden states (norms ranging from 37 to 174), so feeding the wrong combination of states to the fc layer produces completely different activations.

The message also illustrates the importance of understanding data formats at every stage of the pipeline. The assistant had to trace through the training data loader (data.py), the standardization function (standardize_data_v1), the hidden state extraction code, and the SGLang model wrapper to find the mismatch. Each component individually looked correct; it was only when the full chain was examined that the bug became visible.

After implementing the fix — modifying deepseek_v2.py to capture the embedding output and updating the draft model config — the standalone test accuracy jumped from 34.1% to 76.9%, closely matching the training metric. However, the SGLang server still only achieved 54.8 tok/s, suggesting that additional issues remained in the inference pipeline beyond the input format mismatch. The debugging journey was not over, but message 4459 represents the critical breakthrough that unlocked the path forward.