Tracing the Missing Auxiliary Hidden States: Debugging EAGLE-3's Zero Acceptance Rate in SGLang

The Subject Message

In the middle of an intense debugging session, the assistant issued the following command:

ssh root@10.1.230.174 "grep -A20 'def set_eagle3_layers_to_capture' /root/sglang/python/sglang/srt/models/deepseek_v2.py"

The output revealed the critical function definition:

def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None):
    if not self.pp_group.is_last_rank:
        return

    if layer_ids is None:
        self.capture_aux_hidden_states = True
        num_layers = self.config.num_hidden_layers
        self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3]
    else:
        self.capture_aux_hidden_states = True
        # we plus 1 here because in sglang, for the ith layer, it takes...

At first glance, this appears to be a routine code inspection — a developer grepping for a function definition to understand how it works. But in the context of the broader debugging narrative, this single message represents a pivotal moment of discovery. It is the thread that, when pulled, unravels the deepest root cause of why a carefully trained EAGLE-3 draft model with 1.2 billion parameters was producing zero accepted tokens during speculative decoding inference.

The Context: A Long Debugging Journey

To understand why this message matters, we must step back and appreciate the journey that led to it. The assistant and user had been working for days on deploying speculative decoding for the Kimi-K2.5 large language model — a massive Mixture-of-Experts model running across 8 NVIDIA RTX PRO 6000 Blackwell GPUs. The goal was to train an EAGLE-3 draft model: a smaller, auxiliary transformer that predicts multiple future tokens in parallel, which the main model then verifies. If the draft model's predictions are accurate enough, the system achieves significant throughput improvements.

The training pipeline had been a monumental effort. Hidden states were extracted from the target model using a custom server-side patch. Ten thousand training samples were collected. A new EAGLE-3 drafter was trained from scratch using the speculators library. Training logs showed reasonable loss curves and validation metrics. Everything looked promising.

Then came the moment of truth: deploying the trained draft model with SGLang's speculative decoding engine. The server launched successfully, the model loaded, but the acceptance metrics told a devastating story. The log line read: accept_len: 1.05, accept_rate: 0.21. An accept length of 1.05 means that on average, only 0.05 extra tokens beyond the first guaranteed token were accepted per step. The acceptance rate of 0.21 — barely above the 0.20 floor that represents the baseline where only the first token is accepted by construction — meant the draft model was essentially useless.

The First False Lead: Weight Key Names

The assistant initially suspected a weight loading issue. The speculators library saves the decoder layer weights under the key prefix layers.0.*, but SGLang's LlamaForCausalLMEagle3 model expects midlayer.*. A quick fix was written — a Python script to rename the keys in the safetensors file — and the server was restarted. But the acceptance rate barely budged, moving from 0.20 to 0.21. This marginal improvement suggested that the decoder weights were now loading correctly, but the model's predictions were still essentially random.

This was a confounding result. The weights had been verified to have reasonable magnitudes and non-zero values. The shapes matched SGLang's expectations, including the critical QKV projection dimension of 2 * hidden_size = 14336. Yet the draft model was producing predictions no better than random chance.

Shifting the Hypothesis: The Hidden State Pipeline

The assistant then pivoted to a new hypothesis: the problem might not be in the draft model's weights at all, but in the hidden states being fed into it. The EAGLE-3 architecture has a distinctive feature: it doesn't just use the final hidden state from the target model. Instead, it captures hidden states from three different layers of the target model, concatenates them into a single 21504-dimensional vector (3 × 7168), and projects this down to 7168 dimensions through a learned fusion layer called fc. This multi-layer feature fusion is what gives EAGLE-3 its predictive power — the draft model gets a richer representation of the target model's internal state.

The critical question became: was SGLang actually capturing and passing these multi-layer hidden states, or was it only passing a single layer's hidden state (7168 dimensions)? If the latter, the fc fusion layer would be bypassed entirely, and the draft model would receive impoverished features that it was never trained on.

The Subject Message: Tracing the Capture Mechanism

This brings us to the subject message. The assistant is now tracing the hidden state capture mechanism in SGLang's source code. The command greps for set_eagle3_layers_to_capture in the DeepSeekV2 model file — the model architecture that Kimi-K2.5 is based on.

The output reveals the function's logic. When layer_ids is None (the default case), it sets capture_aux_hidden_states = True and selects three layers: layer 2, the middle layer (num_layers // 2), and the third-from-last layer (num_layers - 3). When layer_ids is provided, it uses those specific layer IDs instead.

But the crucial detail is what the function does not show. The output is truncated at a comment: # we plus 1 here because in sglang, for the ith layer, it takes.... More importantly, the function is defined on the DeepseekV2Model class — but the Kimi-K2.5 model might use a different model class that doesn't inherit this method, or might override it incorrectly.

The Deeper Problem: Is set_eagle3_layers_to_capture Even Called?

Looking at the broader context, the assistant had previously discovered (in [msg 3561]) that the model runner calls self.model.set_eagle3_layers_to_capture() only if self.eagle_use_aux_hidden_state is True. This flag is set from the EAGLE configuration, which in turn comes from the server's speculative decoding parameters.

The assistant's earlier investigation (in [msg 3557]) had shown that the draft forward method in eagle_worker.py checks self.eagle_use_aux_hidden_state to decide whether to use hidden_size * 3 as the input dimension for the fusion layer. If this flag is False — or if the target model's set_eagle3_layers_to_capture method is never called because the flag is set too late or the model doesn't have the method — then the auxiliary hidden state capture mechanism is never activated.

Without auxiliary hidden states, the target model only produces a single final-layer hidden state (7168 dimensions). The draft model receives this 7168-dim vector, but its fc layer expects 21504 dimensions. The shape check hidden_states.shape[-1] != embeds.shape[-1] evaluates to 7168 != 7168 — they match! — so the fusion is bypassed entirely. The draft model then operates on impoverished single-layer features, producing predictions that are barely better than random.

The Insight: Why Both Old and New Drafters Failed

This diagnosis explains a previously mysterious observation: both the old vLLM-trained drafter and the new SGLang-trained drafter exhibited identical zero-acceptance behavior. Both were trained on fused multi-layer features (21504 → 7168), but both received single-layer features (7168) at inference time. The training data and the inference data were fundamentally mismatched, guaranteeing failure regardless of the quality of the trained weights.

The weight key rename fix was a red herring — it addressed a real issue, but fixing it only revealed the deeper problem. The decoder weights were now loading correctly, but they were being fed the wrong inputs.

Input Knowledge Required

To fully understand this message, one needs knowledge of several interconnected systems:

  1. The EAGLE-3 architecture: Understanding that the draft model uses auxiliary hidden states from multiple layers of the target model, concatenated and fused through a learned projection layer.
  2. SGLang's speculative decoding infrastructure: How the eagle_worker.py orchestrates the interaction between target and draft models, and how hidden states flow from one to the other.
  3. The DeepSeekV2 model architecture: The base model that Kimi-K2.5 is built upon, and how its layer structure maps to the auxiliary hidden state capture mechanism.
  4. The speculators library: The training framework used to train the draft model, and how it differs from SGLang's inference-time assumptions.
  5. The concept of capture_aux_hidden_states: A mechanism that intercepts hidden states at specific layers during the target model's forward pass, storing them for later use by the draft model.

Output Knowledge Created

This message produces several valuable pieces of knowledge:

  1. The exact implementation of set_eagle3_layers_to_capture in the DeepSeekV2 model, including the default layer selection logic (layers 2, middle, and third-from-last).
  2. Confirmation that the function exists in the DeepSeekV2 model file, meaning it should be available for the Kimi-K2.5 model if properly configured.
  3. The two-branch logic: The function handles both the default case (no layer IDs specified) and the explicit case (specific layer IDs provided), with a comment hinting at a +1 offset convention in SGLang's layer indexing.
  4. A narrowing of the hypothesis space: The problem is not in the weight loading, not in the draft model architecture, and not in the fusion layer dimensions — it is in the upstream hidden state capture mechanism.

Assumptions and Potential Mistakes

The assistant makes several assumptions in this investigation:

  1. That the Kimi-K2.5 model actually uses the DeepSeekV2 model class. If Kimi-K2.5 uses a custom model class that doesn't inherit set_eagle3_layers_to_capture, or overrides it in a way that disables auxiliary hidden state capture, then the entire investigation into the DeepSeekV2 implementation is moot.
  2. That the eagle_use_aux_hidden_state flag is being set correctly. The model runner only calls set_eagle3_layers_to_capture if this flag is True. If the flag is False due to a configuration issue, the auxiliary hidden state mechanism is never activated regardless of the model class.
  3. That the layer indexing convention is consistent. The truncated comment mentions "+1" for SGLang's layer indexing, hinting at a potential off-by-one error between the layer IDs used during training (in speculators) and those used during inference (in SGLang).
  4. That the capture_aux_hidden_states flag on the model is sufficient. Setting this flag to True might not be enough — the model's forward pass must actually implement the logic to store hidden states at the specified layers. If the forward pass doesn't check this flag, the hidden states are never captured.

The Thinking Process

The assistant's reasoning in this message is a textbook example of systematic debugging. The progression is clear:

  1. Observe the symptom: Zero acceptance rate despite trained weights loading correctly.
  2. Formulate a hypothesis: The hidden states passed to the draft model are wrong.
  3. Trace the data pipeline: Follow the flow of hidden states from the target model's forward pass through the capture mechanism to the draft model's input.
  4. Inspect the capture mechanism: Look at set_eagle3_layers_to_capture to understand how layers are selected and how the capture flag is set.
  5. Identify the gap: The function exists in the DeepSeekV2 model, but is it being called? Is the flag being set? Is the Kimi-K2.5 model using this class? The truncated output is particularly telling. The comment about "+1" in layer indexing hints at a subtle but critical detail: SGLang might use 1-based layer indexing while the training framework uses 0-based indexing. If the layer IDs are off by one, the draft model would receive hidden states from the wrong layers, producing features that don't match what it was trained on.

Conclusion

This message, while seemingly a simple code grep, represents the critical transition in a debugging journey. The assistant has moved from surface-level issues (weight key names) to the fundamental architectural question: are the hidden states being captured and passed correctly? The answer to this question will determine whether the entire EAGLE-3 training pipeline needs to be rethought, or whether a simple configuration fix can unlock the draft model's potential.

The set_eagle3_layers_to_capture function is the gatekeeper of the auxiliary hidden state mechanism. If it's not being called — or if it's being called with the wrong parameters — then every downstream component (the fusion layer, the decoder, the lm_head) is operating on corrupted inputs. The trained weights, no matter how good, cannot compensate for fundamentally mismatched features.

The assistant's systematic approach — verifying each link in the chain from weights to hidden states to capture mechanism — demonstrates the kind of disciplined debugging required when deploying complex ML systems. Each hypothesis is tested, each component verified, until the true root cause emerges from the noise.