Tracing the EAGLE-3 Draft Model Architecture: A Deep Dive into ML Debugging

In the middle of a complex debugging session aimed at understanding why a newly trained EAGLE-3 draft model was underperforming during speculative decoding, a seemingly simple command was issued. Message [msg 4406] consists of a single bash tool call that searches for files containing specific code patterns in the speculators Python package. The command and its output are brief, but they sit at a critical juncture in a much larger investigation — one that reveals the painstaking, iterative nature of debugging modern machine learning inference pipelines.

The Context: A Performance Mystery

The broader session had been wrestling with a frustrating problem. The assistant had trained an EAGLE-3 draft model for the Kimi-K2.5 language model on 100,000 samples, achieving a respectable 74.7% validation accuracy during training. Yet when deployed with SGLang's speculative decoding engine, the performance was abysmal: only ~46.7 tokens per second compared to a 90.0 tok/s baseline, with an acceptance length of merely ~1.9 tokens out of 16 draft tokens. The draft model, despite its training accuracy, was simply not predicting well in the inference pipeline.

The assistant had already discovered one critical bug: the --speculative-num-steps 1 flag was silently overriding the --speculative-num-draft-tokens 16 setting, limiting the draft to just 2 tokens due to an internal SGLang constraint when topk=1. But fixing that only made performance worse, revealing that the fundamental issue lay deeper — in how the draft model's forward pass was wired within SGLang versus how it was trained.

The Immediate Trigger: Understanding the Decoder Layer

The message immediately preceding [msg 4406] (msg [msg 4405]) shows the assistant trying to inspect the decoder layer implementation:

grep -n "class Eagle3Decoder\|norm_before_residual\|hidden_norm\|self_attn\|hidden_states.*=.*hidden" 
  ~/ml-env/lib/python3.12/site-packages/speculators/models/eagle3/decoder.py

This command returned empty output — the file didn't exist or didn't contain those patterns. The assistant was stuck. The decoder layer was the missing piece in understanding the full forward pass.

In message [msg 4403], the assistant had reconstructed the architecture from the core.py file:

  1. hidden_states = self.fc(hidden_states) — maps 3×hidden_size → hidden_size
  2. input_embeds = self.embed_tokens(input_ids) — embeds the previous token
  3. hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) — concatenates to 2×hidden_size
  4. Pass through decoder layers with RoPE
  5. logits = self.lm_head(self.norm(hidden_states)) — produce logits But there was a dimensionality puzzle: the decoder layer output should be 2×hidden_size (7168×2 = 14336) based on the concatenation, yet self.norm was an RMSNorm(hidden_size) (7168). How could a norm designed for 7168 dimensions process a 14336-dimensional tensor? The assistant suspected the decoder layer had a special first layer that split and recombined the concatenated embedding and fc outputs, reducing the dimensionality back to hidden_size.

The Search: Message 4406

This is where message [msg 4406] enters. The assistant needed to find where the decoder layer was actually implemented. The previous grep on decoder.py had failed, so the assistant broadened the search:

ssh root@10.1.230.174 'find ~/ml-env/lib/python3.12/site-packages/speculators/ -name "*.py" | 
  xargs grep -l "norm_before_residual\|Eagle3Decoder\|hidden_norm" | head -10'

The output revealed six files containing these patterns:

The Reasoning and Motivation

Why was this search necessary? The assistant was operating at the boundary between two codebases: the speculators package (used for training the draft model) and SGLang (used for inference). The draft model architecture was defined in speculators, but SGLang had its own implementation in llama_eagle3.py. The assistant needed to verify that SGLang's forward pass matched the training architecture exactly. Any discrepancy — even a subtle one in how hidden states were normalized, concatenated, or projected — could explain why a model with 74.7% training accuracy was producing near-random predictions during inference.

The search terms were carefully chosen:

Input Knowledge Required

To understand this message, one needs substantial context about the EAGLE-3 architecture. EAGLE-3 is a speculative decoding framework where a small "draft" model predicts multiple tokens ahead, which are then verified by the full "target" model. The draft model doesn't just take the last hidden state — it takes auxiliary hidden states from multiple layers of the target model (specifically layers 3, 31, and 59 for Kimi-K2.5), concatenates them with the embedding output, projects them down, and passes them through a lightweight transformer decoder.

The assistant also needed to know:

Output Knowledge Created

This message produced a concrete list of files that contain the decoder-related code. The most important discovery was model_definitions.py, which the assistant would immediately read in the next message. This file revealed LlamaDecoderEagle3FirstLayer — a custom decoder layer that inherits from HuggingFace's LlamaDecoderLayer but with modifications to handle the concatenated [embed, fc_output] input format.

The search also confirmed that decoder.py didn't exist as a separate file — the decoder implementation was spread across core.py and model_definitions.py. This was valuable negative knowledge: the assistant could stop looking for a standalone decoder file and focus on these two files instead.

Assumptions and Potential Mistakes

The assistant made several implicit assumptions in this search:

  1. That the decoder implementation would contain one of the three search terms
  2. That the implementation was in the speculators package (not in SGLang or elsewhere)
  3. That norm_before_residual was a distinguishing feature of the decoder These were reasonable assumptions, but they could have missed the decoder if it used different terminology. For instance, if the decoder was implemented purely in terms of generic Transformer blocks without Eagle3-specific naming, this search would have failed. The assistant was fortunate that the custom first layer class was named LlamaDecoderEagle3FirstLayer.

The Broader Narrative

This message, while technically just a file search, represents a crucial step in a debugging process that would ultimately uncover a fundamental wiring mismatch. The assistant would later discover that SGLang was passing the wrong hidden states to the draft model — using cat([layer3, layer31, layer59]) (the three auxiliary hidden states captured by the target model) instead of cat([embed_output, layer3, layer31]) (the embedding output plus the first two auxiliary states). This mismatch explained why the draft model's fc layer was receiving the wrong input format, producing garbage predictions despite correct weights.

The search in message [msg 4406] was the gateway to understanding the decoder architecture well enough to trace the full forward pass and identify this mismatch. Without locating model_definitions.py and understanding how LlamaDecoderEagle3FirstLayer processes the concatenated input, the assistant would have remained stuck on the dimensionality puzzle and might never have found the root cause.

Conclusion

Message [msg 4406] exemplifies the meticulous, iterative nature of debugging complex ML systems. A single grep command, taking seconds to execute, can unlock hours of investigation by revealing where critical code lives. The assistant's methodical approach — tracing the forward pass step by step, verifying each assumption against the actual code, and broadening searches when initial attempts fail — is a masterclass in systematic debugging. The message itself may appear trivial, but it sits at the heart of a narrative about the gap between training accuracy and inference performance, and the painstaking work required to bridge that gap.