The Moment of Truth: Benchmarking a "Fixed" EAGLE-3 Drafter

Introduction

In the long and arduous journey of deploying speculative decoding for the Kimi-K2.5 language model, few moments are as tense as the one captured in message [msg 3551]. After spending hours training a 1.2B-parameter EAGLE-3 draft model on 10,000 samples of hidden states extracted from SGLang, and then debugging a weight key name mismatch that silently dropped all trained weights during model loading, the assistant finally restarts the server with a "fixed" checkpoint and runs the benchmark. This message is the moment of truth — the instant when the assistant discovers whether the fix actually worked.

The message itself is deceptively brief. The assistant notes "Good, it's responding," then launches a single-stream benchmark against the SGLang server running with the newly patched EAGLE-3 draft model. The results show approximately 25.5 tokens per second — a far cry from the 90 tok/s baseline the assistant had achieved earlier without any speculative decoding. The acceptance rate, checked in the immediately following message [msg 3552], remains at a dismal 0.21 (accept_len ~1.05), meaning essentially zero draft tokens are being accepted. The weight key fix addressed a real bug, but it was not the root cause.

The Path to This Moment

To understand why this message exists, one must trace the debugging journey that preceded it. The assistant had trained an EAGLE-3 draft model using hidden states extracted from SGLang (the 10K sample run), and the training metrics looked promising — validation loss plateauing around 6.13 with step-0 accuracy at approximately 74.5%. Yet when the model was deployed for inference, it produced an acceptance rate of exactly 0.20, which with 5 draft tokens meant zero draft tokens were ever accepted.

The assistant's first hypothesis, developed across messages [msg 3537] through [msg 3542], was a weight key name mismatch. The speculators library (used for training) saves the single decoder layer as layers.0.* (e.g., layers.0.self_attn.q_proj.weight, layers.0.mlp.down_proj.weight), but SGLang's LlamaForCausalLMEagle3 class expects these weights under the prefix midlayer.*. The load_weights method in SGLang's eagle model implementation tries to map checkpoint keys to model parameters by checking first the bare key name, then model.{key}. So fc.weight correctly maps to model.fc.weight, but layers.0.hidden_norm.weight fails to match model.midlayer.hidden_norm.weight. The result: the entire decoder layer remains at random initialization, and the draft model produces garbage predictions.

The fix was straightforward: rename all layers.0.* keys to midlayer.* in the safetensors file using a small Python script ([msg 3542]). The assistant then killed the old server, replaced the checkpoint file, and restarted SGLang ([msg 3544]-[msg 3546]). After waiting through the lengthy startup process (over 5 minutes), the server came online.

The Message Itself

The subject message [msg 3551] reads:

Good, it's responding. Now let me run the benchmark and check acceptance rate:

>

[bash] ssh root@10.1.230.174 "/root/ml-env/bin/python3 /tmp/sglang_bench.py --url http://localhost:8000 --mode single --num-requests 5 --max-tokens 256"

>

Warming up... Warm-up done.

>

=== Single-Stream Benchmark === Prompt: Explain the theory of general relativity in detail... Tokens: 256 completion, 18 prompt Time: 10.02s, TPOT: 39.2ms, TPS: 25.5 tok/s

>

Prompt: Write a Python function to implement merge sort wi... Tokens: 256 completion, 20 prompt Time: 10.20s, TPOT: 39.8ms, TPS: 25.1 tok/s

>

Prompt: What are the main causes and effects of climate ch... Tokens: 256 completion, 19 prompt Time: 9.92s, TPOT: 38.7ms, TPS: 25.8 tok/s

>

=== M...

The tone is cautiously optimistic — "Good, it's responding" — but the assistant is clearly not celebrating. The benchmark results tell an immediate story: 25.5 tok/s is catastrophically slow. The baseline without any speculation was 90 tok/s, and even the initial broken speculation run achieved 24.8 tok/s ([msg 3535]). The weight key fix has changed essentially nothing.

Assumptions and Their Failure

The assistant made a reasonable assumption: that the weight key name mismatch was the sole cause of the zero-acceptance behavior. This assumption was grounded in solid evidence — the load_weights code clearly showed that layers.0.* keys would not match midlayer.* parameters, causing the trained weights to be silently dropped. Fixing this mismatch should have loaded the trained decoder layer weights correctly.

But the assumption was incomplete. As the chunk summary reveals, there is a second, deeper issue: the hidden states passed to the draft model are 7168-dimensional (single-layer) instead of the expected 21504-dimensional concatenation of three auxiliary layer hidden states. The fc fusion layer, which projects 21504 → 7168, is never applied because the shape check hidden_states.shape[-1] != embeds.shape[-1] evaluates to 7168 != 7168 → False, bypassing the fusion entirely. The draft model was trained on fused multi-layer features but receives single-layer features at inference time.

This explains why both the old vLLM-trained drafter and the new SGLang-trained drafter exhibit identical zero-acceptance behavior — they both receive the wrong type of input. The root cause is that eagle_use_aux_hidden_state is not properly activated for the KimiK25 model, or the target model's capture_aux_hidden_states mechanism is not producing the multi-layer hidden states that the draft model was trained on.

Input Knowledge Required

To fully understand this message, one needs knowledge of:

  1. Speculative decoding architecture: How EAGLE-3 works as a draft model that predicts multiple tokens per forward pass, with a verification step that accepts or rejects draft tokens.
  2. EAGLE-3's auxiliary hidden state mechanism: The draft model is trained not on the final hidden state of the target model, but on a concatenation of hidden states from multiple intermediate layers (typically 3 layers), fused via an fc projection layer down to the model's hidden dimension. This gives the draft model richer contextual information.
  3. SGLang's model loading conventions: The LlamaForCausalLMEagle3 class uses midlayer as the attribute name for its single decoder layer, while the speculators training library uses layers.0. This naming discrepancy is a cross-framework compatibility issue.
  4. The KimiK25 model architecture: A DeepSeek-V2 derivative with MLA (Multi-head Latent Attention), which may not implement the capture_aux_hidden_states interface that SGLang's EAGLE-3 support expects.
  5. The training pipeline: How hidden states are extracted from the target model, how the draft model is trained on those states, and how the trained weights are saved and loaded.

Output Knowledge Created

This message produces several critical pieces of knowledge:

  1. The weight key fix is insufficient: Despite correctly renaming layers.0.* to midlayer.*, the acceptance rate remains at baseline (zero draft tokens accepted). This eliminates one hypothesis but confirms that a more fundamental issue exists.
  2. The throughput is worse than baseline: At ~25 tok/s with speculation versus 90 tok/s without, the draft model is actively harmful. Even if acceptance were perfect, the overhead of running the draft model and verification would need to be offset by high acceptance rates.
  3. The problem is systematic: Since both the vLLM-trained and SGLang-trained drafters show identical behavior, the issue is not specific to one training run or one framework. It's a fundamental mismatch between how the draft model was trained and how it's being used at inference time.
  4. The debugging must shift focus: The weight key issue was a red herring. The real problem lies in the hidden state pipeline — how SGLang captures and passes hidden states from the target model to the draft model.

The Thinking Process

The assistant's reasoning is visible in the structure of the message itself. The phrase "Good, it's responding" reveals a baseline anxiety — after the long server restart, the first check is simply whether the server is alive. The assistant doesn't jump to conclusions; instead, it methodically runs a benchmark and prepares to check the acceptance rate logs.

The decision to run --mode single --num-requests 5 --max-tokens 256 is deliberate. A single-stream benchmark with short generation lengths (256 tokens) provides a quick signal without consuming excessive time or compute. The assistant knows from previous experience that the acceptance rate is the key metric, and that the server logs will contain accept_len and accept_rate values. The benchmark is a means to generate enough traffic to populate those log entries.

The assistant also shows discipline in not over-interpreting the results prematurely. The message ends with "=== M..." (truncated in the conversation data), suggesting the benchmark output was cut off. The assistant doesn't declare victory or defeat — it simply collects the data and moves to the next step (checking the acceptance logs in [msg 3552]).

Broader Implications

This message illustrates a fundamental challenge in deploying speculative decoding: the training and inference pipelines must be perfectly aligned. A draft model trained on fused multi-layer hidden states will fail catastrophically if the inference server passes single-layer hidden states. The weight key mismatch was a red herring — a real bug, but not the root cause. The deeper issue is that SGLang's EAGLE-3 integration for non-standard model architectures (like KimiK25 with MLA attention) does not properly activate the auxiliary hidden state capture mechanism.

The debugging process here mirrors a pattern common in complex ML systems: you fix one bug, only to discover that the real problem lies deeper. Each fix peels back a layer of abstraction, revealing the next underlying issue. The weight key fix was necessary but not sufficient — it corrected the loading of trained weights, but those weights were trained on features that the inference pipeline never produces.

Conclusion

Message [msg 3551] is a pivotal moment in the EAGLE-3 debugging saga. It represents the culmination of one hypothesis (weight key mismatch) and the beginning of a deeper investigation into the hidden state pipeline. The assistant's methodical approach — fix, restart, benchmark, check logs — demonstrates good engineering discipline even when the results are disappointing. The message itself is brief, but it carries the weight of the entire debugging journey that preceded it and sets the stage for the deeper investigation that follows.

The lesson is clear: in speculative decoding, the training data pipeline and the inference data pipeline must be mirror images of each other. A mismatch at any level — weight naming, hidden state dimensions, feature composition — can render a carefully trained draft model completely useless. The assistant's next steps, revealed in subsequent messages, will involve patching SGLang's hidden state capture mechanism to produce the multi-layer features that the draft model expects.