The Hidden State Wiring Bug: Debugging EAGLE-3 Speculative Decoding from 56.8 tok/s to Root Cause and Beyond

Introduction

In the high-stakes world of large language model inference, every token per second matters. Speculative decoding — a technique where a small "draft" model predicts multiple tokens cheaply and a large "target" model verifies them in parallel — promises transformative throughput gains. But when the numbers don't add up, the debugging journey can reveal deep, unexpected flaws in the pipeline.

This article chronicles a concentrated debugging session spanning dozens of messages in an opencode coding session, where an AI assistant and user worked to deploy an EAGLE-3 draft model for the Kimi-K2.5 1-trillion-parameter Mixture-of-Experts model on 8 NVIDIA RTX PRO 6000 Blackwell GPUs. The draft model had achieved 74.7% validation accuracy during training, yet when deployed with SGLang's speculative decoding, it delivered a paltry 56.8 tokens per second — far below the 90 tok/s non-speculative baseline. What followed was a systematic investigation that uncovered not one but two critical bugs, culminating in the discovery of a silent hidden state wiring mismatch between training and inference that explained the catastrophic performance gap.

The session is a masterclass in systematic ML systems debugging: isolating components, tracing data formats, questioning assumptions, and iterating toward understanding. It demonstrates that in the frontier of LLM inference optimization, the most elusive bugs are often not in the model weights or the training algorithm, but in the silent assumptions that bridge the gap between training and deployment.

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 (TP=8). 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 benchmark at message [msg 4356] told a sobering story: 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 server logs revealed the grim details:

Run 1/5: 2048 tokens in 36.32s = 56.4 tok/s
Run 2/5: 2048 tokens in 33.90s = 60.4 tok/s
Run 3/5: 2048 tokens in 38.80s = 52.8 tok/s
Run 4/5: 2038 tokens in 35.29s = 57.8 tok/s
Run 5/5: 2048 tokens in 35.89s = 57.1 tok/s

With 16 draft tokens configured and most being rejected, the overhead of generating and verifying them was destroying any potential speedup. Something fundamental was wrong.

The First Bug: A Silent Configuration Override

The assistant's first breakthrough came from investigating SGLang's internal parameter handling. A subagent task traced through the server argument code and discovered a critical constraint: when --speculative-eagle-topk 1 is set (the only sensible value for EAGLE-3's chain-based speculation), SGLang silently forces num_draft_tokens = num_steps + 1. The server had been started with --speculative-num-steps 1 (the default), which meant num_draft_tokens was clamped to just 2 — not the 16 the user had explicitly configured. The draft model was only ever proposing one token per step, making the entire speculation mechanism nearly pointless.

The fix was straightforward: set --speculative-num-steps 15 to match the 16 draft tokens. But when the assistant restarted the server and re-ran the benchmark at message [msg 4370], the result was even worse: 46.7 tok/s. The accept length barely budged, hovering around 1.9 tokens. The server logs showed:

accept len: 2.27, accept rate: 0.14
accept len: 1.73, accept rate: 0.11
accept len: 2.08, accept rate: 0.13

This was the moment of crisis — the hypothesis that the num_steps override was the primary bottleneck had been decisively falsified. Something deeper was wrong with how the draft model was being used.

The Diagnostic Pivot: Testing on Training Data

At this critical juncture, the user suggested a diagnostic experiment that would reframe the entire investigation. At message [msg 4375], the user proposed: "Try with a training sample and see if accept rate looks correct. Possible we didn't wire in the model correctly."

This was a classic debugging heuristic: if a model performs well on held-out validation data during training but fails in deployment, test it on training data to distinguish between a generalization problem and an integration problem. If the model also fails on training data, the issue is in how it's loaded or invoked; if it succeeds on training data but fails on new data, the issue is generalization.

The assistant immediately recognized the value of this approach, articulating the diagnostic framework: "If accept rate is high on training data but low on new data, it's a generalization issue. If it's also low on training data, we have a wiring/loading bug." This binary framework would guide the next several hours of work.

Building the Standalone Test

The assistant's next move was a textbook debugging technique: write a standalone test that bypasses the complex serving infrastructure entirely. At message [msg 4379], the assistant wrote test_drafter_standalone.py — a Python script that loaded the draft model weights from safetensors, loaded a training sample's hidden states from disk, manually ran the EAGLE-3 forward pass, and compared the predicted next token against the ground truth.

The script was SCP'd to the remote server and executed at message [msg 4388]. The initial run showed the model weights loading correctly — all expected keys present with the correct shapes and dtypes. But the critical discovery came in the following message ([msg 4389]): 0.0% accuracy on training data. The draft model was completely unable to predict the next token on data it had been explicitly trained on.

This was the smoking gun. The problem was definitively a wiring or loading bug, not a generalization issue. But what kind of wiring bug could cause a model to achieve 0% accuracy on its own training data?

The Vocabulary Mapping Red Herring

The assistant initially suspected a problem with the d2t vocabulary mapping tensor. The draft model uses a reduced vocabulary of 32,000 tokens (down from the target model's 163,840), and the d2t tensor maps draft token IDs to target token IDs using an offset format: target_id = draft_id + d2t[draft_id].

The standalone test showed d2t[:10] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0], meaning the first 10 draft tokens all mapped to target token 0 — clearly broken. The assistant traced through SGLang's llama_eagle3.py and discovered that it correctly converts offsets to direct mappings via hot_token_id = loaded_weight + torch.arange(loaded_weight.shape[0]). The standalone test confirmed this: hot_token_id[:10] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], showing that the first 10 draft tokens map to the first 10 target tokens as expected.

After fixing the d2t handling in the standalone test, accuracy improved to 34.1% — still far below the 74.7% training metric. This investigation, while ultimately a dead end, was valuable because it eliminated a whole class of potential bugs and narrowed the search to the hidden state input format.

The Hidden State Revelation

The assistant's investigation of the training data format revealed a crucial detail. The hidden states directory contained 4 tensors per sample, not the 3 the script initially expected. At message [msg 4385], the assistant examined the tensor norms and discovered a striking pattern:

Tracing Through the Training Pipeline

The answer came from reading the speculators library's data preprocessing code. At message [msg 4462], the assistant read the standardize_data_v1 function in the training library's data.py:

def standardize_data_v1(data: dict[str, Any]) -> dict[str, Any]:
    return {
        "hidden_states": torch.cat(data["hidden_states"][:-1], dim=-1),
        "input_ids": data["input_ids"],
        "verifier_last_hidden_states": data["hidden_states"][-1],
        ...
    }

The critical line was data["hidden_states"][:-1] — Python slice notation that takes all elements except the last one. Since the hidden states list contained [embed_output, layer3_output, layer31_output, layer59_output], the [:-1] slice selected [embed_output, layer3_output, layer31_output]. These three tensors were concatenated along the feature dimension to form the 21,504-dimensional input to the fc layer.

The last hidden state (layer59_output) was stored separately as verifier_last_hidden_states, used for a different purpose in the training loop.

The Root Cause: A Silent Wiring Mismatch

At message [msg 4463], the assistant synthesized the discovery:

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!

The draft model's fc layer had been trained on a specific combination of hidden states: the embedding output (which carries information about the previous token) plus two auxiliary hidden states from layers 3 and 31. But SGLang, configured with eagle_aux_hidden_state_layer_ids = [2, 30, 58], was capturing only the three auxiliary hidden states (layers 3, 31, and 59) and concatenating them — completely omitting the embedding output.

The embedding output is qualitatively different from the layer outputs. It has a norm of ~1 compared to norms of 37–174 for the layer outputs, meaning the fc layer had learned to expect a specific pattern of magnitudes in its input. When it received three large-magnitude layer outputs instead of one small-magnitude embedding plus two large-magnitude layer outputs, the projections were completely wrong.

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 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.
  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 at message [msg 4470]. 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 at message [msg 4491]. 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 at message [msg 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 new hypotheses: 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. At message [msg 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 assistant then investigated whether the TBO (Two-Batch Overlap) path might be bypassing the capture mechanism. When can_run_tbo is True and first_k_dense_replace > normal_start_layer, the normal loop only runs the dense layers, and the remaining MoE layers run through a separate model_forward_maybe_tbo path that doesn't include capture logic. However, the server logs showed disable_overlap_schedule=True, confirming TBO was disabled for speculation.

The Architecture Investigation: Tracing Through the Speculators Library

Before the standalone test could reveal the wiring mismatch, the assistant had to understand the draft model's architecture in detail. This required tracing through the speculators library's core.py and model_definitions.py files to reconstruct the exact forward pass.

At message [msg 4403], the assistant synthesized the forward pass into five steps:

  1. hidden_states = self.fc(hidden_states) — maps 3×hidden_size to hidden_size
  2. input_embeds = self.embed_tokens(input_ids) — embed the previous token
  3. hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) — concatenate to 2×hidden_size
  4. Pass through decoder layers with RoPE
  5. logits = self.lm_head(self.norm(hidden_states)) — project to vocabulary But this raised an immediate question: the concatenation in step 3 produces a 2×hidden_size tensor, yet the RMSNorm in step 5 expects hidden_size. How does the decoder layer reduce the dimension? The answer came from reading the decoder layer source code. The first decoder layer is special — it splits the 2×hidden_size input at the midpoint, applies separate LayerNorms to each half, recombines them for attention, and produces a hidden_size output. This "split-and-recombine" pattern is a non-trivial architectural choice that has profound implications for how hidden states must be prepared before being fed into the draft model. At message [msg 4408], the assistant wrote: "Now I understand the full architecture. The first layer is special." This was followed by a seven-step description of the forward pass that captured the split-and-recombine logic. The assistant then identified the critical next question: "Now the key question: what does SGLang's llama_eagle3.py do vs speculators?" This question drove the comparison that would eventually reveal the hidden state format mismatch. The assistant read the full llama_eagle3.py file from the SGLang source tree and compared it line by line against the speculators implementation. The comparison revealed that SGLang's fc projection only runs if the hidden state dimension doesn't match the embedding dimension — if the hidden states are already 7168-dimensional, the fc layer is skipped entirely. This was a clue that the hidden states arriving at the draft model were not in the expected format.

Tracing the Hidden State Pipeline in SGLang

The assistant also traced how SGLang captures hidden states from the target model. At message [msg 4419], the assistant grepped the DeepSeek V2 model code for layers_to_capture and discovered the mechanism:

self.model.layers_to_capture = [val + 1 for val in layer_ids]

This off-by-one shift means that a layer_ids value of [2, 30, 58] becomes [3, 31, 59] in the actual capture list. The assistant also discovered that the default layer selection formula was [2, num_layers // 2, num_layers - 3] — early, middle, and late layers.

But the critical insight was that SGLang's capture mechanism only captured transformer layer outputs, not the embedding output. There was no mechanism to capture the embedding layer because it's not a transformer layer — it's a separate embedding table that runs before any transformer layers. This is why the training pipeline had to be patched to capture the embedding output separately, and why SGLang's default configuration silently omitted it.

The Broader Lessons

This debugging session illustrates several fundamental principles for deploying speculative decoding systems:

1. Training-inference data format alignment is critical. The most subtle bugs are often not in the model weights or the training algorithm, but in the silent assumptions that bridge the gap between training and deployment. A single Python slice ([:-1]) was the root cause of a catastrophic performance collapse.

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. Configuration parameters can silently override each other. The --speculative-num-steps 1 flag silently overriding --speculative-num-draft-tokens 16 was a hidden dependency that wasted significant debugging time. Understanding the full parameter interaction graph is essential.

5. Embedding outputs are not the same as layer outputs. The embedding layer produces hidden states with fundamentally different statistical properties (norm ~1 vs norms of 37-174 for transformer layers). Confusing them with layer outputs can silently corrupt the input to downstream components.

6. 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.

7. The debugging process is iterative. Each fix revealed new layers of the problem. The num_steps fix exposed the wiring bug, and the wiring fix exposed remaining performance issues. Effective debugging requires the flexibility to abandon hypotheses and pivot when the data doesn't match expectations.

The Unfinished Story

As the segment 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.

The assistant added debug prints to llama_eagle3.py to inspect the actual shapes of hidden states reaching the draft model, and began investigating whether the TBO path might be bypassing the capture mechanism. These investigations were interrupted by an empty user message at [msg 4503], which triggered a context re-establishment.

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 segment 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.

The core discovery — that the draft model's fc layer was trained on cat([embed, layer3, layer31]) while SGLang was passing cat([layer3, layer31, layer59]) — was the essential piece that no amount of hyperparameter tuning or model retraining could have solved. It required tracing the data flow through four separate codebases: the training data extraction pipeline, the speculators library's data preprocessing, the draft model configuration, and SGLang's inference engine. This level of cross-system understanding is what separates effective ML engineers from those who merely tweak parameters and hope for the best.