The Hidden State Wiring Bug: Debugging EAGLE-3 Speculative Decoding from Training to Inference

Introduction

In the complex ecosystem of large language model deployment, speculative decoding has emerged as a powerful technique for accelerating inference without sacrificing quality. The EAGLE-3 algorithm, a sophisticated form of speculative decoding, trains a lightweight "draft" model to predict the target model's next tokens using intermediate hidden states as conditioning signals. But when the hidden state wiring between training and inference diverges — even subtly — a perfectly trained draft model becomes useless.

This article synthesizes a concentrated debugging session spanning dozens of messages, where an AI assistant traced a catastrophic performance failure in an EAGLE-3 speculative decoding deployment from symptom to root cause, implemented a fix, and then confronted the sobering reality that the fix was only partially successful. The session is a masterclass in systematic ML systems debugging: isolating components, tracing data formats, questioning assumptions, and iterating toward understanding.

The Setup: A Promising Model That Failed at Inference

The story begins with a straightforward deployment. The assistant had trained an EAGLE-3 draft model for the Kimi-K2.5 language model — a 1-trillion-parameter Mixture-of-Experts model running on 8 NVIDIA RTX PRO 6000 Blackwell GPUs with tensor parallelism. The draft model was trained on 100,000 samples of hidden states extracted from the target model, achieving a validation accuracy of 74.7%. This was a strong result, suggesting the drafter had learned to predict the target model's next-token distribution with high fidelity.

Yet when deployed with SGLang's speculative decoding engine, the results were catastrophic. The server achieved only ~56.8 tokens per second with an accept length of ~1.6 tokens, compared to a baseline of 90.0 tok/s without speculation. The draft model was barely predicting any tokens correctly — the speculation was essentially useless, and in fact was slowing down inference.

The First Red Herring: Configuration Override

The assistant's first hypothesis was a configuration error. SGLang's --speculative-num-steps 1 flag was silently overriding --speculative-num-draft-tokens 16, reducing the effective draft length to just 2 tokens due to an internal constraint when topk=1. This was a real bug, but fixing it to --speculative-num-steps 15 actually worsened performance to 46.7 tok/s with an accept length of only ~1.9 tokens.

This was the first clue that the problem was deeper than configuration. The draft model was genuinely failing to produce acceptable predictions in the inference environment, despite its strong training metrics. Something fundamental was wrong with how the model was being invoked.

The Breakthrough: Isolating the Draft Model

The assistant's next move was a textbook debugging technique: write a standalone test that bypasses the complex serving infrastructure entirely. The test script (test_drafter_standalone.py) loaded the draft model weights directly, fed them hidden states from the training dataset, and computed next-token prediction accuracy. The goal was to isolate whether the problem was in the draft model's weights (a training issue) or in how SGLang was wiring the model into the inference pipeline (an integration issue).

The initial standalone test results were shocking: only 34.1% accuracy on training data, compared to the 74.7% reported during training. This gap was far too large to attribute to random variation or minor implementation differences. Something fundamental was wrong with how the test was feeding data to the draft model.

Tracing Through the Training Pipeline

The assistant then embarked on a meticulous forensic analysis of the training data pipeline. By examining the speculators library's standardize_data_v1 function, the assistant discovered a critical detail:

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

The training data contained four hidden state tensors per sample, captured in a specific order: [embed_output, layer3, layer31, layer59]. The training pipeline concatenated the first three ([:-1]) — i.e., cat([embed_output, layer3, layer31]) — and used this 21,504-dimensional vector (3 × 7168) as input to the draft model's fully-connected (fc) layer. The fourth tensor (layer59) was stored separately as verifier_last_hidden_states for the verification step.

But the standalone test — and, critically, the SGLang inference pipeline — was concatenating the last three hidden states ([layer3, layer31, layer59]). This meant the fc layer, which had been trained to expect [embed, layer3, layer31] as input, was receiving completely different features. The embedding output — the most informative signal for next-token prediction — was being replaced by the deepest layer's activation (layer59), which represents a much more abstract and task-specific representation.

The assistant confirmed this by examining hidden state norms. The first tensor (hs[0]) had a mean norm of ~1.04 — typical of embeddings, which are relatively smooth. The other tensors had norms in the range of 37–174 — typical of deep layer activations with much larger magnitudes. This confirmed that hs[0] was indeed the embedding output, not a transformer layer output.

The Root Cause: Two Different Conventions

The root cause was now clear. The training pipeline and the SGLang inference pipeline used fundamentally different conventions for selecting which hidden states to feed to the draft model:

The Fix: Modifying SGLang's Source Code

With the root cause identified, the assistant faced a choice: modify the training pipeline to match SGLang's convention (which would require retraining the draft model, taking hours), or modify SGLang's source code to match the training convention (which could be done in minutes). The assistant chose the latter.

Implementing the fix required modifying deepseek_v2.py in SGLang to support capturing the embedding output as an auxiliary hidden state. The challenge was that SGLang's capture mechanism only operated at transformer layer boundaries — it had no concept of capturing the embedding output that exists before any layer runs.

The assistant's reasoning in message <msg id=4473> traced through the SGLang forward loop execution:

  1. "The embedding is computed before the loop." In transformer models, the embedding lookup happens once before the main layer loop.
  2. "At the start of the loop, hidden_states = embed_output, residual = None." The residual variable, which in later layers carries the previous layer's output, is initialized to None.
  3. "So capturing at layer 0 would give hidden_states + residual = embed_output + None." The capture code does aux_hidden_states.append(hidden_states + residual). If residual=None, this would crash. This ruled out the simple approach of using layers_to_capture = 0. Instead, the assistant implemented a more surgical fix:
  4. Added a capture_embedding_for_eagle3 boolean flag to the model class, initialized to False in the constructor.
  5. Modified set_eagle3_layers_to_capture to handle layer_id = -1 as a special value meaning "capture the embedding output." When -1 is detected, the flag is set to True and the negative value is filtered out from the layer capture list (since there's no actual layer -1).
  6. Added capture logic in the forward pass, right after aux_hidden_states = [] is initialized and before the layer loop begins. When capture_embedding_for_eagle3 is True, it clones hidden_states (which at that point is the embedding output) and appends it to the auxiliary states list. The draft model configuration was also updated from [2, 30, 58] to [-1, 2, 30], telling SGLang to capture [embed_output, layer3_out, layer31_out] — matching the training convention exactly.

The Verification: Standalone Test Confirms the Fix

After applying the fix, the assistant ran the standalone test with the corrected input format. The result: 76.9% accuracy, closely matching the 74.7% training metric. This confirmed that the wiring mismatch was the primary cause of the deployment failure. The draft model itself was fine — the pipeline was broken.

This verification was critical. It established that:

The Moment of Reckoning: A Partial Victory

The assistant then restarted the SGLang server with the corrected configuration and ran a benchmark. The result: 54.8 tok/s with an accept length of ~1.8 out of 6 draft tokens — an acceptance rate of roughly 30% per token.

This was a partial victory. The fix had improved performance from the broken 46.7 tok/s, but it was still far below the 90 tok/s baseline without speculation. The acceptance rate of 30% was barely higher than the broken version's ~12% (1.9 out of 16), and far below the expected ~75% based on the standalone test.

The assistant's analysis in message <msg id=4493> was sharp and disciplined:

"Accept length is ~1.8 (accept rate ~0.30 = 1.8/6). This is barely higher than the broken version (~1.9/16). Something is still wrong. The accept rate per drafted token is 30% but should be ~75% based on our offline test."

The assistant immediately generated a new hypothesis: perhaps the hidden_states.clone() operation in the embedding capture patch creates a tensor that is incompatible with tensor parallelism. In SGLang's TP=8 setup, each GPU holds a shard of the hidden states. The clone() operation might produce a tensor that lacks the proper parallelism metadata, causing the draft model to receive corrupted or misaligned data.

Questioning Fundamental Assumptions

As the assistant dug deeper, it began questioning more fundamental assumptions. In message <msg id=4498>, the assistant asked:

"Actually, I realize there may be a more fundamental issue. The hidden state capture during training was done via our custom SGLang dump patch, which captured the hidden states on TP rank 0. But during inference, the hidden states are also on TP rank 0. So the format should match. Let me check something else — perhaps the issue is that the hidden_states variable at the start of the loop already includes the first dense layer's transformation."

This pivot reveals a sophisticated debugging mindset. Rather than tweaking parameters or trying random fixes, the assistant went back to first principles: "What exactly is the tensor at the point where we capture it?" The assistant read the source code to verify that the embedding output hadn't already passed through some transformation before reaching the capture point.

The Broader Lessons

This debugging session illustrates several fundamental lessons for ML engineering:

1. Training and inference pipelines must be exactly aligned. The EAGLE-3 architecture is particularly sensitive to hidden state wiring because it operates on the target model's internal representations rather than token-level predictions. Small differences in which hidden states are captured, how they are concatenated, or how positions are shifted can completely invalidate the draft model's predictions.

2. Silent correctness bugs are the most dangerous. The hidden state mismatch produced no error messages, no dimension mismatches, no crashes. The model silently computed wrong predictions because both the correct and incorrect input formats had the same shape (21,504 dimensions). Only by comparing accuracy against training metrics could the bug be detected.

3. Isolate before you investigate. The standalone test was the breakthrough that made everything else possible. By eliminating SGLang's complex inference pipeline from the equation, the assistant could confirm that the draft model itself was correct and narrow the search to the interface between the model and the serving infrastructure.

4. A fix that should work but doesn't reveals deeper issues. The assistant's fix was correct — it aligned the input format with the training convention. But the persistent performance gap (30% acceptance rate vs 75% expected) showed that the input format mismatch was only part of the problem. The assistant's willingness to question its own assumptions and go back to the source code is what separates effective debugging from random experimentation.

5. Document and contextualize. Throughout the session, the assistant maintained a running record of hypotheses tested, results obtained, and conclusions drawn. This discipline paid off when the empty user message at <msg id=4503> triggered a comprehensive context re-establishment that became the definitive reference for all subsequent work.

The Unfinished Story

As the chunk concludes, the debugging effort remains unresolved. The hidden state wiring fix improved performance from 46.7 tok/s to 54.8 tok/s, but the acceptance rate of 30% remains far below the expected 75%. The assistant has generated new hypotheses about tensor parallelism interactions and the exact semantics of the hidden_states variable at the capture point, but these have yet to be tested.

This unfinished ending is itself instructive. In real-world ML engineering, bugs rarely have single, clean causes. The input format mismatch was real and needed fixing, but it wasn't the whole story. The assistant's methodical approach — measure, compare, hypothesize, test, and iterate — provides a template for tackling the deeper issues that remain.

Conclusion

The debugging session captured in this chunk is a microcosm of the challenges in deploying custom-trained speculative decoding models. The gap between offline accuracy (76.9%) and online acceptance rate (30%) highlights how subtle differences between training and inference pipelines can catastrophically degrade performance. The wiring mismatch that took hours to discover — a simple [:-1] slicing difference — cost the assistant an entire debugging cycle.

But the session also demonstrates a robust methodology for tackling such challenges: isolate the component, test against known data, compare against training metrics, trace discrepancies to their root cause, implement surgical fixes, and when those fixes fall short, question fundamental assumptions and return to first principles. In the complex world of ML systems debugging, where a single off-by-one error can cost days of computation, this disciplined process is the difference between spinning wheels and making progress.