Tracing the Hidden State Flow: Implementing EAGLE-3 Extraction in SGLang

Introduction

In the complex world of speculative decoding for large language models, few tasks are as delicate as extracting intermediate hidden states from a running inference server. Message <msg id=3277> captures a pivotal moment in this process: the assistant is deep in the bowels of SGLang's codebase, tracing how aux_hidden_states — the intermediate layer outputs needed for EAGLE-3 training — flow through the logits_processor module. This single grep command, seemingly mundane, represents the culmination of hours of debugging, a server hang, a killed process, and a strategic pivot that would ultimately unlock the ability to train a dramatically better speculative drafter.

The Broader Context: Why Hidden States Matter

To understand the significance of this message, we must step back. The assistant had been working on deploying the Kimi-K2.5 model (a 547GB, 8-GPU behemoth) using SGLang, and had been pursuing EAGLE-3 speculative decoding — a technique where a lightweight "drafter" model predicts multiple tokens per forward pass, accelerating inference. The first attempt at EAGLE-3 had failed: the drafter achieved only ~15% acceptance rate and 0.66x throughput, making speculative decoding slower than base inference.

The root cause was training data quality. The initial drafter was trained on hidden states extracted via vLLM, but the assistant had identified that SGLang's different attention implementation (triton vs. flashinfer) and model execution path produced different hidden state distributions. The solution was to extract hidden states directly from SGLang during inference, ensuring the drafter learned from the exact same representations it would encounter during deployment.

But SGLang was not designed for this. Unlike vLLM, which has explicit hooks for speculative decoding frameworks like speculators, SGLang's internals are tightly coupled — the model forward pass is woven into the serving infrastructure via ForwardBatch objects, pipeline parallelism, and CUDA graph capture. Extracting hidden states required patching the model at exactly the right point in the forward pass.

The Message: A Surgical Investigation

The message itself is brief but dense:

Now let me see how logits_processor handles aux_hidden_states — it likely stores them in forward_batch:

``bash ssh root@10.1.230.174 "grep -n 'aux_hidden_states' /root/sglang/python/sglang/srt/layers/logits_processor.py | head -20" ``

292: aux_hidden_states: Optional[torch.Tensor] = None, 323: aux_hidden_states, 330: aux_hidden_states, 403: aux_hidden_states: Optional[torch.Tensor], 417: if aux_hidden_states is not None: 418: aux_pruned_states = [hidden for hidden in aux_hidden_states] 445: if aux_hidden_states is not None: 446: aux_pruned_states = [hidden[last_index] for hidden in aux_hidden_states] 541: aux_hidden_states: Optional...

This single command reveals the assistant's investigative methodology. Rather than reading documentation (which likely doesn't exist for this internal mechanism), the assistant is tracing the code path empirically — searching for every reference to aux_hidden_states in the logits processor to understand how intermediate hidden states are received, transformed, and potentially stored.

The Reasoning: Why the Logits Processor?

The assistant's comment — "it likely stores them in forward_batch" — reveals a critical hypothesis. In SGLang's architecture, the ForwardBatch object is the central data structure that carries per-request state through the model forward pass. If aux_hidden_states were stored in forward_batch, they would be accessible after the forward pass completed, allowing the assistant to save them to disk from a higher-level wrapper without modifying the core model code.

The grep results confirm the hypothesis partially. The logits_processor does receive aux_hidden_states as a parameter (lines 292, 403, 541) and processes them — pruning to the last token position for draft model training (lines 417-418, 445-446). But critically, the grep doesn't show aux_hidden_states being stored back into forward_batch. This means the hidden states flow through the forward pass as a separate tensor, are consumed by the logits processor, and are then discarded.

This discovery has profound implications for the extraction strategy. If the hidden states are ephemeral — created during the forward pass and destroyed afterward — the assistant cannot simply read them from forward_batch after inference. Instead, the patch must intercept them during the forward pass, before the logits processor consumes them.

Input Knowledge Required

To understand this message, several layers of knowledge are required:

SGLang Architecture: The reader must understand that SGLang uses a ForwardBatch object to carry per-request state through the model, that it supports pipeline parallelism across multiple GPUs, and that the logits_processor is the final stage that converts hidden states to token logits.

EAGLE-3 Training Pipeline: The assistant is working within the speculators framework, which requires hidden states from three intermediate layers (typically layers 3, 31, and 59 of a 60-layer model) plus the final layer's hidden state. These are the "auxiliary hidden states" — intermediate representations that the drafter model learns to predict.

The Previous Failure: The grep is informed by the knowledge that the previous EAGLE-3 drafter failed due to poor acceptance rates, and that the root cause was likely a mismatch between training data distribution and deployment distribution.

SGLang's Model Code Structure: The assistant had already traced the aux_hidden_states flow through deepseek_v2.py (the model file) in previous messages, finding that they are captured at layers [3, 31, 59] and returned alongside the final hidden state. The logits processor is the next link in the chain.

Output Knowledge Created

This message produces several concrete insights:

The logits processor receives aux_hidden_states as a direct parameter, not via forward_batch. This confirms the assistant's suspicion about where to intercept them.
The hidden states are pruned to last-token positions (lines 445-446: aux_pruned_states = [hidden[last_index] for hidden in aux_hidden_states]). This is important because EAGLE-3 training only needs the last token's hidden state from each layer, not the full sequence.
The hidden states are consumed and discarded within the logits processor — there's no evidence they're stored back to forward_batch or any other persistent structure. This means the patch must capture them at the point of creation, not after the fact.
The code path is confirmed to exist and be active — the capture_aux_hidden_states flag (seen in earlier grep results) controls whether the model captures these states at all. When enabled, they flow through the entire forward pass to the logits processor.

Assumptions and Their Validity

The assistant operates under several assumptions in this message:

Assumption 1: The logits processor is the right place to intercept. This is reasonable — if hidden states are stored in forward_batch after logits processing, the assistant could save them without modifying the model forward pass. However, the grep results disprove this assumption, forcing a different approach.

Assumption 2: The grep results are complete. The head -20 limit means the assistant only sees the first 20 matches. There could be additional references beyond line 541 that show storage into forward_batch. This is a practical trade-off — reading the entire file would be more thorough but slower.

Assumption 3: The SGLang installation is consistent. The assistant is grepping a file at /root/sglang/python/sglang/srt/layers/logits_processor.py, assuming this is the exact code being executed. If the running server uses a different installation path or a modified version, the grep results could be misleading.

Assumption 4: The aux_hidden_states mechanism is the correct path. The assistant assumes that the capture_aux_hidden_states flag (seen in the model code) is the mechanism used for EAGLE-3-style extraction. This is a well-founded assumption based on the code structure, but it's worth noting that there could be alternative mechanisms.

The Thinking Process Visible in the Message

The assistant's thinking is revealed through the structure of the investigation. This message is part of a systematic trace through the code:

Previous messages: The assistant traced aux_hidden_states through deepseek_v2.py, finding they're captured at layers [3, 31, 59] and returned from the model forward pass.
This message: The assistant follows the chain to the next consumer — the logits processor — to understand how the states are used and whether they persist.
Next steps (implied): The assistant would need to either (a) patch the model forward to save hidden states before returning them, (b) patch the logits processor to save them before consumption, or (c) intercept at the CausalLM level where the model output is received. The comment "it likely stores them in forward_batch" is a hypothesis being tested. The assistant expects to find evidence of storage but instead finds evidence of consumption. This negative result is valuable — it rules out the simplest approach and forces a more invasive patch strategy.

The Broader Significance

This message, while technically a simple grep command, represents a critical juncture in the EAGLE-3 training pipeline. The assistant had already:

Killed a hung SGLang server (flashinfer backend deadlocked on SM120)
Restarted with triton attention + NCCL tuning (which would later achieve 90 tok/s)
Begun designing the hidden state extraction approach The investigation of logits_processor was the final piece of reconnaissance before implementing the server-side patch. The discovery that hidden states are ephemeral — consumed and discarded by the logits processor — directly shaped the patch design. The assistant would need to modify the model's forward method to save hidden states to disk (as binary .pt files in /dev/shm/) before they reach the logits processor. This approach — Approach C in the assistant's planning — would prove successful. The full 10K-sample extraction completed with zero errors, producing 17.3M tokens of hidden states (924 GB) and enabling the training of a new EAGLE-3 drafter from scratch. That drafter achieved ~74% step 0 accuracy, ~64% step 1 conditional accuracy, and ~56% step 2 conditional accuracy — dramatically better than the previous broken drafter's 25% acceptance rate.

Conclusion

Message <msg id=3277> is a masterclass in systematic debugging. On its surface, it's a simple grep command — a developer searching for a variable name in a source file. But in context, it's the decisive moment where the assistant confirms the architecture of SGLang's hidden state flow and commits to a particular patching strategy. The grep results, showing that aux_hidden_states are consumed and discarded within the logits processor, ruled out the simpler approach of reading them from forward_batch after inference and forced the development of the more invasive server-side patch that would ultimately succeed.

This kind of surgical code tracing — following a data flow through multiple files, forming and testing hypotheses about where data persists, and adapting strategy based on empirical findings — is the essence of working with complex, undocumented systems. The assistant's methodical approach, building on hours of prior investigation, demonstrates how deep understanding of a system's internals can unlock capabilities that the system was never designed to support.