The Pivot That Uncovered a Wiring Bug: Testing EAGLE-3 on Training Samples

"Try with a training sample and see if accept rate looks correct. Possible we didn't wire in the model correctly"

This seemingly simple suggestion, made by the user at message index 4374, represents a critical turning point in a deep debugging session. The context is a sophisticated speculative decoding pipeline for the Kimi-K2.5 language model, where an EAGLE-3 draft model had been trained to 74.7% validation accuracy but was delivering catastrophic inference performance—46.7 tokens per second versus a 90 tok/s baseline, with an average accept length of only ~1.9 out of 16 drafted tokens. The assistant had spent several rounds chasing configuration issues, discovering that --speculative-num-steps 1 was silently overriding --speculative-num-draft-tokens 16 to just 2 tokens, and then finding that even after fixing this to --speculative-num-steps 15, performance actually worsened. The draft model was rejecting tokens at an alarming rate despite its impressive training metrics, and the assistant was running out of hypotheses.

The Reasoning Behind the Message

The user's message is deceptively brief, but it encodes a sophisticated debugging strategy. The core insight is a classic machine learning diagnostic: disentangle model quality from integration correctness. The assistant had been operating under the assumption that the trained draft model was good (74.7% accuracy) and the problem lay in how SGLang was using it—hence the focus on configuration parameters like speculative-num-steps, speculative-eagle-topk, and their interactions. But the user recognized a more fundamental possibility: perhaps the model was never correctly integrated into the inference pipeline at all.

The phrase "Possible we didn't wire in the model correctly" is the key. It acknowledges that between the training pipeline (which used the speculators library) and the inference pipeline (which used SGLang's custom EAGLE-3 implementation), there could be a mismatch in how hidden states are formatted, how the vocabulary mapping is applied, or how the draft model's forward pass is invoked. Testing on training data—where the correct next tokens are known—would immediately reveal whether the draft model's predictions are sensible when given the exact same inputs it saw during training.

Assumptions and Their Implications

The message rests on several implicit assumptions. First, that training data is accessible and can be fed through the inference pipeline—a non-trivial assumption given that the training data was stored as JSONL files with tokenized sequences and hidden states in separate sharded directories. Second, that the training accuracy of 74.7% was genuine and not an artifact of data leakage or overfitting to a specific hidden state format. Third, that if the model performs well on training data but poorly on live inference, the problem is generalization (distribution shift); if it performs poorly on both, the problem is a wiring bug.

The user's assumption that a wiring bug was possible turned out to be prescient. The assistant had been assuming that because SGLang's llama_eagle3.py correctly handled the d2t offset-to-direct mapping (line 241-243: self.hot_token_id = loaded_weight + torch.arange(loaded_weight.shape[0])), and because the eagle_config specified the correct layer IDs [2, 30, 58], the pipeline was correctly wired. But as the subsequent investigation revealed, there was a subtle mismatch in which hidden states were being concatenated as input to the draft model's fully-connected projection layer.

Input Knowledge Required

To understand this message, one needs substantial background knowledge. The reader must understand speculative decoding—the concept of a lightweight "draft" model that proposes multiple candidate tokens which a larger "target" model then verifies in parallel. They need to know about EAGLE-3 specifically, which uses auxiliary hidden states captured from intermediate layers of the target model as conditioning input to the draft model. The architecture involves concatenating three hidden states (each of dimension 7168) into a 21504-dimensional vector, projecting it down to 7168 via a learned fc layer, and then passing it through a transformer decoder layer alongside token embeddings.

One must also understand the training/inference asymmetry that was the root cause: during training, the speculators library used cat([embed_output, layer3, layer31]) as input to the fc layer (taking the embedding output as the first of four hidden states), while SGLang's inference pipeline was passing cat([layer3, layer31, layer59]) (the three auxiliary hidden states captured by the target model, omitting the embedding). This meant the draft model was receiving completely different inputs than it was trained on.

Output Knowledge Created

This message catalyzed an extensive investigation that produced several concrete outcomes. First, the assistant wrote a standalone test script (test_drafter_standalone.py) that loaded the draft model weights directly and ran inference on a training sample, bypassing SGLang entirely. This test initially showed 0.0% accuracy—a shocking result that confirmed the user's suspicion of a wiring bug. Second, the investigation traced through SGLang's codebase to verify that the d2t vocabulary mapping was correctly converted from offset format to direct mapping (confirmed at llama_eagle3.py:241-243). Third, it traced the entire hidden state capture pipeline through deepseek_v2.py, logits_processor.py, and eagle_worker.py to understand exactly how auxiliary hidden states were collected and concatenated.

The most critical output was the discovery that the training pipeline used cat([embed_output, layer3, layer31]) while SGLang was passing cat([layer3, layer31, layer59]). This mismatch meant the fc layer—which expects 3×7168 = 21504 dimensions—was receiving a different combination of hidden states than it was trained on. The fix required modifying deepseek_v2.py to capture the embedding output when layer_id=-1 is specified, and updating the draft model config from [2, 30, 58] to [-1, 2, 30].

The Thinking Process Revealed

The user's message reveals a debugging philosophy that prioritizes ground truth validation over parameter tuning. Rather than continuing to tweak SGLang configuration flags (which had already consumed several rounds of effort), the user stepped back to ask the most fundamental question: does the model work at all when given inputs it understands? This is the equivalent of checking whether a car's engine runs on a test stand before debugging the transmission.

The brevity of the message is also telling. It's a single sentence with two clauses: an action ("Try with a training sample") and a hypothesis ("Possible we didn't wire in the model correctly"). The user isn't prescribing how to test—they leave the implementation to the assistant. They're providing the strategic insight while trusting the assistant to execute the tactical investigation. This division of labor is characteristic of effective human-AI collaboration: the human provides domain intuition and high-level debugging strategy, while the AI handles the detailed code tracing and experimentation.

Broader Significance

This message exemplifies a pattern that recurs across machine learning engineering: when a model performs well on validation metrics but poorly in deployment, the first question should always be whether the deployment pipeline correctly reproduces the training-time computation. The gap between training accuracy (74.7%) and deployment accept rate (~12% per drafted token) was far too large to be explained by distribution shift alone. The user's instinct to test on training data was the correct diagnostic maneuver, and it ultimately uncovered a bug that would have been nearly impossible to find through configuration tweaking alone.

The story also illustrates the hidden complexity of speculative decoding systems. An EAGLE-3 pipeline involves not just loading model weights, but correctly orchestrating hidden state capture from specific target model layers, concatenating them in the right order, applying the correct vocabulary mapping (offset vs. direct), and ensuring the draft model's forward pass matches the training-time computation graph. A single mismatch anywhere in this chain can collapse performance from 74.7% accuracy to near-zero, as the standalone test initially showed.