Unraveling the EAGLE-3 Draft Model Architecture: A Deep Dive into Hidden State Wiring

In the complex world of speculative decoding for large language models, few things are as frustrating as a model that trains beautifully but fails at inference. This article examines a pivotal moment in a debugging session — message [msg 4402] — where an AI assistant, deep in the trenches of diagnosing a poorly performing EAGLE-3 draft model, makes a critical architectural discovery by reading the source code of the speculators library. This message represents the turning point where confusion begins to crystallize into understanding.

The Context: A Model That Should Work, But Doesn't

The session leading up to this message had been a rollercoaster. The team had trained an EAGLE-3 draft model for the Kimi-K2.5 large language model, achieving a respectable 74.7% validation accuracy during training. Yet when deployed with SGLang for speculative decoding, the model produced abysmal results: an accept length of only ~1.6 tokens and throughput of ~56.8 tok/s compared to a 90.0 tok/s baseline — meaning the speculative decoding was actually slower than running the base model directly.

The assistant had already eliminated several possible causes. It confirmed that SGLang correctly handles the d2t vocabulary mapping (converting offset-based mappings to direct token IDs). It had written standalone tests that bypassed SGLang entirely to isolate the draft model's behavior. But those tests initially showed 0% accuracy, leading to a wild goose chase investigating the vocabulary mapping format. After resolving that confusion, the standalone test showed only 34.1% accuracy — far below the 74.7% reported during training.

Something was fundamentally wrong with how the draft model was being used at inference time. The assistant needed to understand the exact architecture of the EAGLE-3 draft model to find the mismatch.

The Subject Message: Reading the Blueprints

Message [msg 4402] is the moment the assistant stops guessing and starts reading the actual source code. The message consists of two parts: a textual observation about a critical finding, followed by a bash command that reads the forward method of the EAGLE-3 draft model from the installed speculators package.

The assistant writes:

Key finding at line 357: hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) — the model concatenates [embed, fc_output] to create 2*hidden_size input before passing to the decoder layer. And at line 220: modified_config.hidden_size = modified_config.hidden_size * 2 — the rotary embedding uses 2*hidden_size.

This is the first glimpse of the actual architecture. The assistant has identified two crucial facts:

The decoder layer input is a concatenation of the embedding output and the fc projection output, not just the fc output alone. This means the decoder operates on a 2*hidden_size vector.
The rotary embedding (RoPE) is configured with double the hidden size, which makes sense given the concatenated input. The assistant then runs sed -n "283,400p" on the speculators core.py file to read the full forward method signature and implementation. This is the equivalent of pulling out the architectural blueprints — the raw source code that defines exactly how the draft model processes hidden states.

Why This Message Matters: The Architecture Revelation

To understand why this message is so important, we need to appreciate what the assistant was up against. The EAGLE-3 architecture is a sophisticated speculative decoding framework where a lightweight "draft" model predicts multiple tokens in parallel, which are then verified by the full "target" model. The draft model doesn't process raw tokens — it receives pre-computed hidden states from the target model at specific layers, compresses them through a fully-connected (fc) layer, and then runs a small transformer to predict draft tokens.

The critical question was: what exactly goes into the fc layer, and what comes out?

The training pipeline had been set up to extract hidden states from the target model at four specific points: the embedding output (layer -1 conceptually), and the outputs of layers 3, 31, and 59. The training data stored these as a list of four tensors: [embed_output, layer3_output, layer31_output, layer59_output].

The standardize_data_v1 function in the speculators library (which the assistant would discover in subsequent messages) processes this list by taking [:-1] — all but the last element — and concatenating them along the last dimension. This means the fc layer was trained on cat([embed_output, layer3_output, layer31_output]), a 21504-dimensional vector (3 × 7168).

But SGLang, in its inference pipeline, was capturing hidden states from the target model at layers 3, 31, and 59 — the three auxiliary hidden states — and concatenating those. This means SGLang was feeding cat([layer3_output, layer31_output, layer59_output]) into the fc layer, completely omitting the embedding output and including layer 59 instead.

This is a fundamental wiring mismatch. The fc layer learned to project a specific combination of hidden states (embedding + early layer + middle layer) into a compressed representation. At inference, it was receiving a completely different combination (three middle-to-late layers). No wonder the predictions were poor.

The Thinking Process Visible in the Message

The message reveals a methodical investigative approach. The assistant doesn't just run code blindly — it reads source code to understand the architecture from first principles. The key insight comes from recognizing that the forward method signature expects hidden_states with shape [1, total_seq_len, 3 * hidden_size], and then traces how those hidden states flow through the model.

The assistant's reasoning shows an understanding that the architecture has a specific, non-obvious structure: the fc layer compresses 3× hidden_size down to hidden_size, but then the result is concatenated with the embedding before entering the decoder. This means the decoder sees both the raw embedding information and the compressed auxiliary state information simultaneously — a design that allows the draft model to leverage both the current token's embedding and the contextual information from the target model's hidden states.

Assumptions and Knowledge Required

To fully grasp this message, one needs to understand several layers of context:

EAGLE-3 speculative decoding: A technique where a small draft model predicts multiple future tokens using hidden states extracted from the main (target) model. The draft model is trained to predict tokens that the target model would accept.
Hidden state capture: During inference, the target model's internal representations at specific layers are "captured" and passed to the draft model. The choice of which layers to capture is critical.
The fc (fully-connected) layer: A linear projection that compresses the concatenated hidden states from 3× hidden_size down to hidden_size, serving as the interface between the target model's representations and the draft model's transformer.
The speculators library: A Python package that provides the EAGLE-3 model implementation, including the Eagle3DraftModel class with its custom forward method.
Rotary Position Embedding (RoPE): A position encoding scheme used in transformer models. The fact that it operates on 2× hidden_size is a consequence of the concatenated input structure.

Output Knowledge Created

This message creates several pieces of critical knowledge:

The draft model's decoder operates on 2× hidden_size, not hidden_size, because it concatenates the embedding with the fc output.
The fc layer output is only half of the decoder input — the other half comes from the embedding of the current token.
The RoPE configuration is doubled to accommodate the 2× hidden_size input.
The exact architecture is now documented via the source code reading, providing a reference for debugging.

The Broader Significance

This message is a textbook example of how debugging complex ML systems often requires going beyond surface-level investigation. The assistant couldn't just look at accuracy numbers or benchmark throughput — it had to read the actual source code of both the training library (speculators) and the inference engine (SGLang) to find the mismatch.

The discovery made here — that the fc layer input format differs between training and inference — would ultimately be confirmed in subsequent messages. The assistant would go on to modify deepseek_v2.py in SGLang to capture the embedding output (using layer_id=-1), update the draft model config from [2, 30, 58] to [-1, 2, 30], and restart the server. After this fix, the standalone test accuracy jumped to 76.9%, matching the training metrics.

However, even after fixing this wiring issue, the production throughput only improved to 54.8 tok/s — still far below the 90.0 baseline. This suggests that while the hidden state input format was one problem, other issues remained in the inference pipeline. The message thus represents a necessary but not sufficient step toward solving the overall performance problem.

Conclusion

Message [msg 4402] captures the moment when a frustrating debugging session transitions from confusion to understanding. By reading the source code of the EAGLE-3 draft model implementation, the assistant discovers the true architecture of the model — revealing that the decoder input is a concatenation of the embedding and the fc output, operating on 2× hidden_size. This architectural insight would prove crucial for identifying the hidden state wiring mismatch between training and inference, ultimately leading to a fix that restored the draft model's predictive accuracy, even if the full performance story remained unresolved.