The Hidden State Hunt: Tracing EAGLE-3's Auxiliary Capture Pipeline in SGLang

In the midst of a deep debugging session on EAGLE-3 speculative decoding performance, a single message marks a critical pivot point in the investigation. The message, <msg id=4412>, reads in its entirety:

[assistant] Now let me look at how the target model captures and concatenates hidden states: [bash] ssh root@10.1.230.174 'sed -n "850,900p" ~/sglang/python/sglang/srt/speculative/eagle_worker.py' -1, ) else: mamba_steps_to_track = None

>

self.target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify( accepted_steps=accepted_steps, mamba_track_indices=batch.mamba_track_indices, mamba_steps_to_track=mamba_steps_to_track, model=self.target_worker.model_runner.model, )

>

def forward_draft_extend( self, batch: ScheduleBatch, hidden_states:...

At first glance, this appears to be a routine grep — the assistant reads a section of source code and gets back output that shows the tail end of one function and the signature of another. But to understand why this message was written, and why it represents a crucial investigative step, we must trace the reasoning chain that led to it.

The Debugging Arc: From Performance Crisis to Wiring Mismatch

The broader context is a multi-session effort to deploy a custom-trained EAGLE-3 draft model for the Kimi-K2.5 large language model using SGLang. The draft model had been trained to 74.7% validation accuracy — promising numbers on paper — yet when deployed with speculative decoding, it achieved only 54.8 tok/s against a 90 tok/s baseline. Something was fundamentally wrong: the draft model was barely being accepted, with an average accept length of ~1.8 tokens out of 6 draft tokens.

The assistant had already eliminated several possible causes. Earlier messages ([msg 4390] through [msg 4396]) confirmed that the vocabulary mapping (d2t → hot_token_id) was handled correctly by SGLang — the offset-to-direct conversion at llama_eagle3.py:241-243 was verified to produce the correct target token IDs. The standalone test written in [msg 4396] revealed a critical wiring mismatch: the training pipeline concatenated [embed_output, layer3, layer31] as input to the draft model's fc layer, but SGLang was passing [layer3, layer31, layer59] — three auxiliary hidden states from the target model, without the embedding output. This explained why the standalone test achieved 76.9% accuracy (matching training metrics) while the SGLang-integrated version failed.

The Question That Drove This Message

With the input format mismatch identified and patched (changing the draft model config from [2, 30, 58] to [-1, 2, 30] to include the embedding layer), the assistant had restarted the server. Performance improved marginally — from 46.7 tok/s to 54.8 tok/s — but remained far below the 90 tok/s baseline. This suggested that while the input format fix was necessary, it was not sufficient. Additional issues lurked in the inference pipeline.

Message <msg id=4412> represents the assistant's next logical step: trace how the target model actually captures and concatenates hidden states during SGLang's speculative decoding pipeline. The reasoning is: if the hidden states passed to the draft model are still wrong even after the config fix, perhaps the target model's capture mechanism itself is flawed — or perhaps the hidden states are being assembled differently than expected.

What the Assistant Was Looking For

The assistant's stated intent — "look at how the target model captures and concatenates hidden states" — reveals a specific investigative hypothesis. In the previous message ([msg 4409]), the assistant had discovered a critical piece of SGLang's LlamaModel.forward():

hidden_states = forward_batch.spec_info.hidden_states
if hidden_states.shape[-1] != embeds.shape[-1]:
    hidden_states = self.fc(hidden_states)

The fc projection — which maps the concatenated auxiliary hidden states (3 × 7168 = 21504 dimensions) down to the model's hidden size (7168) — only runs if the last dimension doesn't match the embed dimension. If the hidden states arriving at the draft model are already 7168-dimensional (single-layer output rather than concatenated three-layer output), the fc layer would be silently skipped, and the draft model would receive raw single-layer hidden states instead of the fused multi-layer representation it was trained on.

This is a subtle but devastating bug: the draft model expects a 21504-dimensional input (three concatenated hidden states from different layers), but if SGLang passes only a single layer's 7168-dimensional output, the fc layer sees matching dimensions and does nothing. The draft model then operates on impoverished features.

The Target Model's Hidden State Capture Mechanism

To understand what actually gets passed, the assistant needed to trace the full pipeline: from the target model's forward pass, through hidden state capture, to the concatenation logic, and finally to the draft model's input. The eagle_worker.py file is the orchestrator — it manages both the target model (the large base model) and the draft model (the small EAGLE-3 model), coordinating the speculative decoding loop.

The specific lines the assistant queried (850-900 of eagle_worker.py) sit at a critical junction. The output shows update_mamba_state_after_mtp_verify — a function that handles state cleanup after verification — and the beginning of forward_draft_extend, which is the method that feeds hidden states into the draft model. The forward_draft_extend method signature (batch: ScheduleBatch, hidden_states:...) is where the assembled hidden states arrive.

What the assistant was hunting for specifically was the code path where layers_to_capture is used. Earlier messages had traced this: in deepseek_v2.py, the model checks if i in self.layers_to_capture during its forward pass and appends hidden states to aux_hidden_states. The model returns (hidden_states, aux_hidden_states) when layers_to_capture is non-empty. The eagle worker then presumably concatenates these captured hidden states before passing them to the draft model.

The Assumptions at Play

Several assumptions underpin this investigation:

  1. The hidden state capture mechanism is properly activated for KimiK25. The assistant had confirmed ([msg 4423]) that KimiK25.set_eagle3_layers_to_capture delegates to self.language_model.set_eagle3_layers_to_capture, which is DeepseekV3ForCausalLM. But the question remained: was the delegation chain actually working? Was the layers_to_capture list being populated correctly?
  2. The concatenation order matches training. The training pipeline concatenated [embed_output, layer3, layer31] — but SGLang's default for EAGLE-3 might concatenate in a different order. The config change from [2, 30, 58] to [-1, 2, 30] was meant to fix this, but if the actual capture mechanism ignores -1 or handles it differently, the fix would be incomplete.
  3. The fc layer dimension check is the only guard. The assistant assumed that if the dimensions match, the fc layer is skipped — meaning if the hidden states are already 7168-dimensional, they pass through unmodified. But what if the hidden states are 21504-dimensional but the fc layer has a different weight shape than expected? Or what if there's a separate path that bypasses the dimension check entirely?

Input Knowledge Required

To understand this message, one needs:

Output Knowledge Created

This message, combined with the subsequent investigation ([msg 4413] through [msg 4426]), produced several critical findings:

  1. The capture path is confirmed: The eagle worker calls forward_draft_extend(hidden_states=logits_output.hidden_states) where logits_output.hidden_states comes from the target model's verification step with CaptureHiddenMode.FULL.
  2. The delegation chain works: KimiK25 → DeepseekV3ForCausalLM → DeepseekV2's layers_to_capture mechanism is functional, with set_eagle3_layers_to_capture being called at model runner initialization ([msg 4426]).
  3. The layers_to_capture default: When called without arguments, set_eagle3_layers_to_capture defaults to [2, num_layers // 2, num_layers - 3] — approximately the early, middle, and late layers. But the custom config [-1, 2, 30] should override this.
  4. The embedding capture question: Layer index -1 is meant to capture the embedding output (before any transformer layers). The assistant would later need to verify that deepseek_v2.py's forward pass actually handles layer_id=-1 by capturing the initial hidden states before the layer loop.

The Thinking Process Visible in This Message

What makes this message fascinating is what it reveals about the assistant's investigative methodology. The assistant is systematically peeling back layers of abstraction:

  1. Start from the symptom: Poor speculative decoding performance despite correct training accuracy.
  2. Eliminate simple causes: Vocab mapping? Verified correct. Config parameters? Fixed from --speculative-num-steps 1 to 15.
  3. Isolate the draft model: Standalone test confirmed the model works correctly in isolation.
  4. Identify the wiring mismatch: Training vs. SGLang use different hidden state compositions.
  5. Fix and retest: Performance improves but doesn't reach baseline.
  6. Dig deeper: If the config fix was correct but performance is still poor, the hidden state capture mechanism itself might be broken — not just the config that specifies which layers to use. This message is step 6 in that chain. The assistant isn't guessing randomly — they're following the data flow from the target model's forward pass, through hidden state capture, to the draft model's input. Each grep and sed command targets a specific link in that chain. The choice to look at lines 850-900 specifically is telling. The assistant already knew from earlier exploration that forward_draft_extend is the entry point for hidden states into the draft model. By examining the code just before that function — the update_mamba_state_after_mtp_verify call — the assistant is checking whether there's any preprocessing or state manipulation that might affect the hidden states before they reach the draft model.

The Broader Significance

This message, while brief, represents a textbook example of debugging a complex ML inference pipeline. The assistant is operating at the intersection of: