The Standalone Test That Exposed EAGLE-3's Hidden State Wiring Bug

Introduction

In the complex world of speculative decoding for large language models, few things are more frustrating than a draft model that achieves 74.7% validation accuracy during training yet delivers only ~12% acceptance rate in production. This discrepancy — a textbook symptom of a deployment wiring bug — became the central mystery of Segment 31 in an ambitious effort to deploy EAGLE-3 speculative decoding with the Kimi-K2.5 model on SGLang. Message 4388 represents the pivotal moment when the assistant ran a standalone test script designed to isolate the draft model from the SGLang inference engine, and the subsequent analysis (in the following message) would reveal a critical input format mismatch that explained the catastrophic performance gap.

The Message

The subject message is a bash command execution on a remote server running Ubuntu 24.04 with 8 NVIDIA RTX PRO 6000 Blackwell GPUs:

[assistant] [bash] ssh root@10.1.230.174 '~/ml-env/bin/python3 /tmp/test_drafter_standalone.py 2>&1'
Vocab mapping: target_vocab=163840, draft_vocab=32000

Draft model weights:
  d2t: torch.Size([32000]) torch.int64
  embed_tokens.weight: torch.Size([163840, 7168]) torch.bfloat16
  fc.weight: torch.Size([7168, 21504]) torch.bfloat16
  lm_head.weight: torch.Size([32000, 7168]) torch.bfloat16
  midlayer.hidden_norm.weight: torch.Size([7168]) torch.bfloat16
  midlayer.input_layernorm.weight: torch.Size([7168]) torch.bfloat16
  midlayer.mlp.down_proj.weight: torch.Size([7168, 18432]) torch.bfloat16...

The output is truncated in the conversation data, showing only the beginning of the weight listing. The script successfully loaded the draft model from its safetensors file and printed the shapes and dtypes of each weight tensor. This seemingly mundane output was the foundation for the explosive discovery in the very next message (msg 4389): 0.0% accuracy on training data.

Why This Message Was Written

The motivation for this message traces back to a series of increasingly desperate debugging steps. The SGLang server had been configured with --speculative-algorithm EAGLE3 and a draft model trained on 100,000 samples, yet the speculative decoding throughput was abysmal. Initial benchmarks showed ~56.8 tok/s with 16 draft tokens, compared to a 90 tok/s baseline without speculation — meaning the drafter was actually slowing down generation.

The user's instruction at [msg 4374] was the catalyst: "Try with a training sample and see if accept rate looks correct. Possible we didn't wire in the model correctly." This was a brilliant diagnostic suggestion. If the draft model performs well on training data but poorly on new data, the problem is generalization. If it performs poorly even on training data, the problem is a wiring or loading bug — a much more tractable issue.

The assistant responded by writing a standalone test script ([msg 4379]) that would load the draft model weights, load a training sample's hidden states from disk, run the EAGLE-3 forward pass manually, and compare the predicted next token against the ground truth. This bypassed SGLang entirely, removing the inference engine as a variable. The script was scp'd to the server ([msg 4380]) and run for the first time ([msg 4381]), but that initial run revealed that the hidden states directory contained 4 tensors per sample — not the 3 the script expected. After inspecting the data format ([msg 4384] and [msg 4385]), the assistant discovered that the training pipeline stored 4 hidden states: [embed_output, layer_3_output, layer_31_output, layer_59_output]. The script was rewritten to handle this format ([msg 4386]), scp'd again ([msg 4387]), and executed in message 4388.

The Debugging Chain: Assumptions, Mistakes, and Discoveries

Several assumptions and potential mistakes are embedded in the context leading to this message:

Assumption 1: The training data format matched what SGLang expected. The assistant initially assumed the hidden states file contained 3 tensors (the auxiliary hidden states from layers 3, 31, and 59). In reality, the training pipeline stored 4 tensors, with the first being the embedding output. This mismatch between training and inference data formats would prove to be the root cause of the poor acceptance rate.

Assumption 2: The d2t mapping was a direct index mapping. The d2t tensor (draft-to-target vocabulary mapping) stores offsets — target_id = draft_id + d2t[draft_id] — not direct target IDs. This was discovered in the following message ([msg 4389]) when the assistant investigated why the first 20 values of d2t were all zero, which would mean all those draft tokens map to target token 0 under a direct interpretation. The offset format is a memory-efficient encoding, but it requires the inference engine to apply the offset correctly.

Assumption 3: SGLang interpreted d2t the same way as the training pipeline. Subsequent investigation ([msg 4391] through [msg 4393]) revealed that SGLang uses hot_token_id as a direct mapping (hot_token_id[draft_id] = target_id), not an offset mapping. This means the same d2t tensor was being interpreted differently by training and inference — a classic wiring bug.

Mistake: The initial test script didn't account for the embedding layer. The first version of test_drafter_standalone.py (written at [msg 4379]) assumed 3 hidden states, but the data had 4. The script was corrected before message 4388, but this highlights how easy it is to mismatch the data pipeline.

Mistake: The SGLang configuration used --speculative-num-steps 1 with --speculative-num-draft-tokens 16. Earlier in the segment ([msg 4360]), the assistant discovered that SGLang silently overrides num_draft_tokens to num_steps + 1 when topk=1, meaning only 2 draft tokens were actually being used despite the 16 requested. This was fixed by setting --speculative-num-steps 15, but even then performance was poor (46.7 tok/s), pointing to a deeper issue with the draft model itself.

Input Knowledge Required

To understand message 4388, one needs knowledge of:

  1. EAGLE-3 speculative decoding architecture: The draft model consists of an embedding layer (embed_tokens.weight), a fully-connected projection layer (fc.weight), a single transformer layer (midlayer), a language modeling head (lm_head.weight), and vocabulary mapping tensors (d2t, t2d). The fc.weight shape [7168, 21504] indicates it projects from a 7168-dimensional hidden state to a 21504-dimensional concatenated input (3 × 7168).
  2. The training data pipeline: Hidden states are extracted from the target model (Kimi-K2.5) at specific layers and stored alongside input IDs and loss masks. The 4 hidden states per sample correspond to the embedding output and the outputs of layers 3, 31, and 59.
  3. SGLang's speculative decoding framework: The inference engine captures auxiliary hidden states from the target model during decoding and feeds them to the draft model. The configuration of which layers to capture is critical.
  4. Vocabulary mapping: The Kimi-K2.5 target model has a vocabulary of 163,840 tokens, while the draft model uses a smaller vocabulary of 32,000 tokens. The d2t and t2d tensors map between these vocabularies.
  5. The hardware context: 8 NVIDIA RTX PRO 6000 Blackwell GPUs with tensor parallelism (tp-size 8), running on Ubuntu 24.04 with CUDA 13.1.

Output Knowledge Created

Message 4388 produced several important pieces of knowledge:

  1. Confirmation that the draft model weights load correctly: All expected keys are present with the correct shapes and dtypes. The fc.weight shape [7168, 21504] confirms the projection layer expects a 21504-dimensional input (3 concatenated 7168-dim hidden states).
  2. The vocabulary mapping tensors are present: d2t (32000 int64) and embed_tokens.weight (163840 × 7168) are both loaded, confirming the draft model has the full target vocabulary embedding.
  3. The model architecture is as expected: hidden_size=7168, intermediate_size=18432, draft_vocab_size=32000, matching the training configuration.
  4. The test harness works: The script successfully loads weights, loads hidden states from disk, and runs the forward pass. This creates a foundation for the accuracy measurement in the next message.
  5. The stage is set for the critical discovery: The output itself doesn't reveal the 0.0% accuracy — that comes in message 4389. But message 4388 is the necessary prerequisite: it proves the model loads correctly, so any accuracy failure must be in the data format or the forward pass logic, not in weight loading.

The Thinking Process

The assistant's reasoning is visible across the chain of messages leading to 4388. The key insight was the user's suggestion to test on training data — this reframed the problem from "why is the model generalizing poorly?" to "is the model even wired correctly?" The assistant's response at [msg 4375] shows this reasoning explicitly: "If accept rate is high on training data but low on new data, it's a generalization issue. If it's also low on training data, we have a wiring/loading bug."

The discovery of 4 hidden states instead of 3 ([msg 4385]) was a crucial moment. The assistant noted the norms: hs[0] had norm ~1.0 (embedding output), while hs[1], hs[2], and hs[3] had norms of 37, 174, and 111 (layer outputs). This immediately suggested that the first hidden state was qualitatively different — the embedding layer output — and the other three were the auxiliary hidden states from layers 3, 31, and 59.

The decision to rewrite the test script to handle 4 hidden states shows adaptive debugging: rather than forcing the data into the expected format, the assistant adjusted the test to match the actual data format. This flexibility was essential for the subsequent discovery.

Conclusion

Message 4388 appears, on its surface, to be a routine diagnostic command — just another bash execution in a long debugging session. But it represents the culmination of a carefully reasoned investigation strategy. The user's suggestion to test on training data, combined with the assistant's methodical approach of isolating the draft model from SGLang, created the conditions for discovering the hidden state input format mismatch that was silently sabotaging speculative decoding performance.

The message itself is the hinge point: before it, the team had been chasing performance issues through configuration parameters (num-steps, num-draft-tokens, topk). After it, the focus shifted to the fundamental wiring between training and inference — the hidden state concatenation order, the vocabulary mapping interpretation, and the embedding layer capture mechanism. This reorientation would eventually lead to the fix documented in the chunk summary: modifying deepseek_v2.py to capture the embedding output when layer_id=-1 is specified, and updating the draft model config from [2, 30, 58] to [-1, 2, 30].

In the broader narrative of this coding session, message 4388 exemplifies a critical debugging principle: when a model works in training but fails in deployment, suspect the interface between them before questioning the model itself.