The Hidden State Wiring Fix: Aligning EAGLE-3 Training and Inference

# Message 4484 — Config update for EAGLE-3 draft model
ssh root@10.1.230.174 '~/ml-env/bin/python3 -c "
import json
with open(\"/data/eagle3/output_100k_sglang/4/config.json\") as f:
    config = json.load(f)
config[\"eagle_config\"][\"eagle_aux_hidden_state_layer_ids\"] = [-1, 2, 30]
with open(\"/data/eagle3/output_100k_sglang/4/config.json\", \"w\") as f:
    json.dump(config, f, indent=2)
print(\"Updated config:\")
print(json.dumps(config[\"eagle_config\"], indent=2))
"
Updated config:
{
  "eagle_aux_hidden_state_layer_ids": [
    -1,
    2,
    30
  ],
  "use_aux_hidden_state": true
}

At first glance, message 4484 appears to be a trivial operation: a three-line Python script that reads a JSON config file, changes one array value from [2, 30, 58] to [-1, 2, 30], and writes it back. But this single edit is the culmination of an intricate debugging odyssey that spanned multiple hours and involved tracing data pipelines across two separate codebases, writing standalone test harnesses, and patching the SGLang inference engine's model forward pass. The change of a single integer — replacing 58 with -1 — represents the resolution of a fundamental mismatch between how the EAGLE-3 draft model was trained and how it was being served at inference time.

The Debugging Journey That Led Here

To understand why this message matters, we must trace the reasoning that produced it. The assistant had been working on deploying an EAGLE-3 speculative decoding drafter for the Kimi-K2.5 model using SGLang. The draft model had been trained on 100,000 samples and achieved a respectable 74.7% validation accuracy. Yet when deployed, it was producing abysmal results: an acceptance rate so low that the speculative decoding pipeline was actually slower than the base model, achieving only ~56.8 tokens per second against a 90.0 tok/s baseline.

The assistant's investigation proceeded through several stages. First, it discovered that the --speculative-num-steps 1 flag was silently overriding the --speculative-num-draft-tokens 16 setting, reducing the effective draft chain to just 2 tokens. After fixing this to --speculative-num-steps 15, performance actually worsened to 46.7 tok/s with an accept length of only ~1.9 — confirming that the draft model itself was producing poor predictions, despite the 74.7% training accuracy.

To isolate the problem, the assistant wrote a standalone test that loaded the draft model directly and fed it hidden states from the training data, bypassing SGLang entirely. This test initially showed poor accuracy too, but then the assistant made a crucial discovery by examining the training data pipeline.

The Root Cause: A Silent Wiring Mismatch

The breakthrough came when the assistant inspected the standardize_data_v1 function in the speculators training library ([msg 4462]). This function processes the raw hidden state data captured during training extraction. The training extraction captured four hidden states per sample: the embedding output, the output after layer 3, the output after layer 31, and the output after layer 59. These were stored as a list of four tensors: [embed_output, layer3_output, layer31_output, layer59_output].

The critical line was torch.cat(data["hidden_states"][:-1], dim=-1) — the Python slice [:-1] takes all elements except the last one. So the training pipeline concatenated [embed_output, layer3_output, layer31_output] into a 21504-dimensional vector (3 × 7168) that served as input to the draft model's fully-connected (fc) projection layer.

Meanwhile, SGLang's EAGLE-3 pipeline was configured with eagle_aux_hidden_state_layer_ids = [2, 30, 58]. The SGLang codebase adds 1 to each of these IDs to compute layers_to_capture = [3, 31, 59], capturing the hidden states before layers 3, 31, and 59 execute. This produced [layer3_output, layer31_output, layer59_output] — a completely different set of three hidden states.

As the assistant noted in [msg 4470]: "The model is fine — the issue is purely a wiring mismatch." The draft model's fc layer had been trained to expect [embed, layer3, layer31] but at inference it was receiving [layer3, layer31, layer59]. The model was being fed garbage, and its predictions were correspondingly garbage.

The Two-Part Fix

The solution required changes in two places. First, the SGLang inference engine needed to be modified to support capturing the embedding output — the hidden state before any transformer layers run. The assistant identified that SGLang's deepseek_v2.py model file had a set_eagle3_layers_to_capture method that translated layer IDs using a +1 convention. To capture the embedding output, the assistant needed layers_to_capture to include 0 (capture before layer 0 = embedding output). With the +1 convention, this meant the config needed to specify -1 as a layer ID.

The assistant made three surgical edits to deepseek_v2.py ([msg 4476], [msg 4479], [msg 4481]):

  1. Modified set_eagle3_layers_to_capture to handle -1 specially, setting a capture_embedding_for_eagle3 flag and filtering out negative IDs from the layers_to_capture list.
  2. Added self.capture_embedding_for_eagle3 = False initialization in the model's __init__.
  3. Added code in the forward pass to capture the embedding output (a clone of hidden_states when residual is None) right after initializing the aux_hidden_states list, before the layer loop begins. The second part of the fix — and the content of message 4484 — was updating the draft model's configuration file to use [-1, 2, 30] instead of [2, 30, 58]. This tells SGLang to capture the embedding output (via -1), the output after layer 2 (which becomes layer 3 after the +1 translation), and the output after layer 30 (which becomes layer 31). This exactly matches the training data format: [embed_output, layer3_output, layer31_output].

Why This Message Matters

Message 4484 is the moment where the assistant commits the config change that makes the entire pipeline coherent. Without this change, all the SGLang code modifications would be useless — the model would still be configured to capture the wrong hidden states. The config file is the bridge between the training pipeline and the inference engine, and this single edit ensures that both sides agree on what data the draft model expects.

The choice of -1 is particularly interesting. It's a sentinel value — a convention that means "capture the embedding output." The assistant had to verify that this would work correctly by examining the SGLang forward pass code ([msg 4472]). At the start of the layer loop, hidden_states contains the embedding output and residual is None. The capture code normally does aux_hidden_states.append(hidden_states + residual), which would crash with None. So the assistant added a guard: hidden_states.clone() if residual is None else hidden_states + residual. This is a pragmatic fix that handles the edge case without restructuring the capture logic.

Assumptions and Knowledge Required

To understand this message, one needs substantial background knowledge. The reader must understand the EAGLE-3 speculative decoding architecture, where a lightweight draft model predicts tokens using hidden states extracted from the target model. They need to know that the draft model's fc layer concatenates multiple hidden states from different layers to form a richer representation. They must understand the SGLang inference engine's convention of capturing hidden states before a layer executes (hence the +1 offset) versus the training pipeline's convention of capturing states after layers execute.

The assistant made several assumptions that proved correct: that the training data's hidden state order was [embed, layer3, layer31, layer59], that standardize_data_v1 used [:-1] to select the first three, and that modifying SGLang to support a -1 layer ID was feasible without breaking other functionality. One assumption that required verification was whether residual=None at the embedding stage would cause a crash — the assistant checked this and added a guard.

The Output Knowledge Created

This message creates a corrected configuration that aligns the training and inference data pipelines. It transforms the draft model from a non-functional state (producing random predictions because it was fed the wrong inputs) to a properly functioning one where the standalone test showed 76.9% accuracy matching training metrics. The config change is the final piece that, combined with the SGLang code modifications, enables the EAGLE-3 speculative decoding pipeline to work correctly.

The broader lesson is a cautionary tale about data pipeline consistency in machine learning systems. When training and inference are implemented in different codebases — here, the speculators training library and the SGLang inference engine — silent mismatches in data format can render a perfectly trained model useless at deployment time. The fix required tracing the data flow through both systems, understanding the conventions each one uses, and finding the translation layer that bridges them. In this case, the translation was a single integer change in a JSON config file, but discovering that change required deep investigation across the entire system.