The Data-Driven Debug: How Empirical Norm Analysis Uncovered the Root Cause of Poor EAGLE-3 Speculative Decoding

Introduction

In the complex world of speculative decoding for large language models, debugging poor performance often feels like searching for a needle in a haystack. When an EAGLE-3 draft model achieves only a 19% token acceptance rate instead of the expected 50%+, the possible causes are numerous: incorrect hidden state wiring, mismatched training data formats, architectural incompatibilities, or simply undertrained weights. Message [msg 4567] captures a pivotal moment in this debugging journey—a moment where the assistant shifts from hypothesis-driven reasoning to empirical, data-driven investigation by running a custom script to compare training data statistics against inference-time hidden state distributions.

The Message: Copying and Executing a Norm Analysis Script

The message itself is deceptively simple—a single bash command that copies a Python script to a remote server and executes it:

[assistant] [bash] scp /home/theuser/glm-kimi-sm120-rtx6000bw/eagle3-train/check_hs_norms.py root@10.1.230.174:/tmp/check_hs_norms.py && ssh root@10.1.230.174 '~/ml-env/bin/python3 /tmp/check_hs_norms.py'
input_ids shape: torch.Size([1281])
Number of hidden state layers: 4
  hs[0] (embed): shape=torch.Size([1281, 7168]), dtype=torch.bfloat16
    per-token norms: mean=1.0383, first5=[1.6596885919570923, 1.4272500276565552, 1.1371872425079346, 1.0681484937667847, 0.9908252954483032]
    first token, first5 values: [0.029541015625, -0.0113525390625, -0.0169677734375, -0.0179443359375, -0.018310546875]
  hs[1] (layer3): shape=torch.Size([1281, 7168]), dtype=torch.bfloat16
    per-token norms: mean=37...

The output is truncated mid-stream, but the critical data is already visible: the training data contains hidden states for 4 layers (embedding, layer 3, layer 31, and layer 59), each with 7168 dimensions, stored in bfloat16 precision. The per-token norm statistics reveal the relative scale of each layer's contribution to the final concatenated hidden state vector.

The Context: A Week of Debugging EAGLE-3 Performance

To understand why this message was written, we must step back into the broader narrative. The assistant had spent the better part of a week deploying and optimizing an EAGLE-3 speculative decoding system for the Kimi-K2.5 model on a machine with 8 RTX PRO 6000 Blackwell GPUs. The journey had been arduous: building flash-attn from source, resolving CUDA version conflicts, training the draft model on 100K samples, and finally deploying it with SGLang's speculative decoding engine.

The previous segment ([msg 4552] through [msg 4562]) had been a deep dive into why the draft model was performing so poorly. The assistant had confirmed that the hidden state shapes were correct (3×7168=21504 dimensions during the extend phase), that the embedding capture was working during verify, and that the autoregressive draft steps correctly used the draft model's own hidden states. Yet the acceptance rate stubbornly remained at 19%.

A critical realization had occurred in [msg 4562]: the assistant noticed that the absolute scale of hidden state norms differed significantly between training and inference. During training, the layer 31 hidden states had per-token norms around 727, while during inference verify they were around 31.3. This discrepancy raised a fundamental question: were the hidden state distributions actually compatible with what the draft model was trained on?

The Reasoning: Why Norm Analysis Matters for EAGLE-3

The EAGLE-3 architecture works by concatenating hidden states from multiple layers of the target model to form a rich input representation for the draft model. In this case, the draft model receives a 21504-dimensional vector formed by concatenating the outputs of layers 2, 30, and 58 (captured at layers 3, 31, and 59 due to SGLang's indexing). The draft model's first layer is a fully-connected (fc) projection that maps this 21504-dim vector down to 7168 dimensions.

The critical insight is that the fc layer was trained on specific hidden state distributions. If the inference-time hidden states have significantly different norms, scales, or statistical properties than the training-time hidden states, the fc projection will produce poor-quality inputs, leading to low acceptance rates.

The assistant's reasoning chain in [msg 4562] reveals this concern explicitly:

"The relative magnitudes are similar (embed and layer3 small, layer31 dominant), but the absolute scales differ a LOT. This is expected since different token content produces different magnitudes."

The assistant was cautiously optimistic—the relative magnitudes were similar, suggesting the wiring was correct—but the scale difference warranted investigation. The check_hs_norms.py script was designed to settle this question definitively by computing per-token norm statistics directly from the training data.

The Script: What check_hs_norms.py Actually Does

While the script's source code is not visible in the message, its output reveals its structure. It loads a training data sample (1281 tokens), extracts the hidden states for all 4 layers (embedding, layer 3, layer 31, layer 59), and computes per-token L2 norms. The output shows:

Assumptions Made by the Assistant

Several assumptions underpin this message:

1. The training data format is correct. The assistant assumes that the hidden states stored during the training data extraction phase accurately represent what the draft model should expect during inference. This is a foundational assumption—if the training data itself was corrupted or incorrectly formatted, no amount of inference-side debugging would help.

2. A single training sample is representative. The script loads one sample (1281 tokens) and computes statistics from it. The assistant implicitly assumes that this single sample's statistics are representative of the entire 100K-sample dataset. Given that the training data was generated from diverse prompts, this assumption is reasonable but not guaranteed.

3. The norm statistics are a sufficient diagnostic. The assistant assumes that comparing per-token norms between training and inference will reveal whether the hidden state distributions match. While norms are a useful aggregate statistic, they can mask more subtle distributional differences that might affect the fc layer's performance.

4. The remote server is in a consistent state. The script is copied and executed on a machine that had been running an SGLang server with the target model loaded. The assistant assumes this doesn't interfere with the analysis.

Mistakes and Incorrect Assumptions

The most significant potential mistake is the assumption that the training data hidden states are the "ground truth" that inference should match. In reality, the assistant had previously made a critical error in this area: in earlier segments, the assistant had incorrectly added embedding capture (with layer_id=-1) to the hidden state extraction, believing the training data included the embedding output. It was later discovered that the training data had never captured the embedding output—the original config [2, 30, 58] was correct all along, and the embedding capture was an erroneous "fix" that actually broke the wiring.

This history is crucial context for understanding [msg 4567]. The assistant is now being extra cautious, empirically verifying that the training data statistics match what the inference pipeline produces. The previous mistake—assuming the embedding was needed—had cost significant debugging time. Now the assistant is letting the data speak for itself rather than relying on assumptions about what the architecture "should" do.

Another subtle issue is that the script only computes norms, not full distributional statistics. Two distributions can have identical L2 norms but completely different value patterns (e.g., one sparse and one dense). The fc layer, being a linear transformation, is sensitive to the full distribution, not just the norm. A more thorough analysis might include histograms, mean/variance per dimension, or PCA projections.

Input Knowledge Required

To fully understand this message, one needs:

1. EAGLE-3 architecture knowledge. Understanding that the draft model receives concatenated hidden states from multiple target model layers, and that a fully-connected projection layer maps this concatenated vector to the draft model's hidden dimension.

2. SGLang's speculative decoding internals. Knowledge of how SGLang's CaptureHiddenMode works, how the eagle worker extracts hidden states during verify, and how accepted_indices selects which hidden states to propagate to the next draft cycle.

3. The training data pipeline. Understanding that hidden states were extracted from the target model using a custom patch that captured outputs at specific layers (3, 31, 59), and that standardize_data_v1 concatenated them as cat([layer3_out, layer31_out, layer59_out]).

4. The debugging history. Knowing that the assistant had previously made an incorrect "fix" by adding embedding capture, and had only recently realized this was wrong.

5. Python/PyTorch data analysis. Basic familiarity with tensor shapes, dtypes, and norm computations.

Output Knowledge Created

This message produces several concrete outputs:

1. Quantitative training data statistics. The per-token norm means for each hidden state layer provide a baseline for comparison. The embedding layer has norms around 1.0, layer 3 around 37, and the truncated output would show layer 31 and layer 59 norms.

2. Validation that the training data contains 4 layers. The output confirms that the training data stores hidden states for the embedding, layer 3, layer 31, and layer 59—consistent with the standardize_data_v1 concatenation order.

3. Confirmation of bfloat16 precision. The hidden states are stored in bfloat16, matching the model's native precision and confirming no precision mismatch between training and inference.

4. A reusable analysis script. The check_hs_norms.py script can be re-run on different samples or compared against inference-time captures to validate distributional compatibility.

The Thinking Process: A Methodical Debugging Approach

What makes this message particularly instructive is the thinking process it reveals. The assistant is not randomly trying fixes or guessing at the problem. Instead, it follows a systematic, hypothesis-driven approach:

  1. Observe the symptom: 19% acceptance rate is far below expectations.
  2. Verify the wiring: Confirm that hidden state shapes and concatenation order match between training and inference.
  3. Identify a discrepancy: Notice that norm scales differ between training and inference.
  4. Formulate a hypothesis: Perhaps the distributional mismatch is causing the fc layer to produce poor projections.
  5. Design an experiment: Write a script to compute training data norms and compare them against inference captures.
  6. Execute the experiment: Copy the script to the remote server and run it (this message).
  7. Analyze the results: Use the output to determine whether the distributions are compatible. This is textbook scientific debugging—measure, don't guess. The assistant resists the temptation to make another "fix" based on intuition (which had already led to the incorrect embedding capture change) and instead gathers empirical evidence.

The Broader Significance

While [msg 4567] is a single message in a long conversation, it represents a turning point in the debugging process. The norm analysis would eventually lead the assistant to discover that the hidden state wiring was actually correct all along—the original config [2, 30, 58] was right, and the embedding capture "fix" had been the real problem. After reverting that change, the acceptance rate jumped from 19% to 47%.

This message also exemplifies a broader lesson in ML engineering: when debugging complex systems, always verify your assumptions against real data. The assistant's willingness to write a custom analysis script, run it on actual training data, and compare the results against inference-time measurements is what ultimately led to the correct diagnosis. In a field where "it should work" often collides with "it doesn't work," the ability to instrument, measure, and compare is the most valuable debugging skill.

The truncated output in this message—showing only the first two layers' statistics—is almost poetic in its incompleteness. The full picture would only emerge after more analysis, more profiling, and more systematic optimization. But this message captures the moment when the assistant chose data over intuition, measurement over assumption, and empirical evidence over guesswork. That choice is what made the subsequent breakthrough possible.