The Discovery That Saved Days: How SGLang's Built-in Hidden States API Changed the Course of EAGLE-3 Training

Introduction

In the middle of a sprawling, multi-session effort to deploy speculative decoding for the massive Kimi-K2.5 language model, a single message from the AI assistant in an opencode coding session represents a pivotal turning point. Message [msg 3283] captures the moment when the assistant realized that SGLang—the serving framework it was already using—had a built-in feature for returning hidden states during inference, completely obviating the need for a complex, error-prone server-side patch that had been the focus of the preceding hours.

This article examines that message in depth: the reasoning that led to the discovery, the assumptions it overturned, the technical knowledge it required and produced, and how it reshaped the trajectory of the entire project. For a reader unfamiliar with the conversation, this is the story of how careful investigation of existing infrastructure can save enormous engineering effort.

The Broader Context: Why Hidden States Matter

To understand the significance of this message, one must understand what the assistant was trying to accomplish. The project involved deploying the Kimi-K2.5 model—a massive 547GB Mixture-of-Experts (MoE) language model with Multi-head Latent Attention (MLA)—across 8 NVIDIA RTX PRO 6000 Blackwell GPUs. The ultimate goal was to implement EAGLE-3 speculative decoding, a technique where a smaller "drafter" model predicts the large model's outputs, enabling faster inference.

Training an EAGLE-3 drafter requires extracting intermediate hidden states from the large model during inference. Specifically, the drafter needs hidden states from specific layers (layers 3, 31, and 59 in this case) to learn how to predict the model's behavior. This extraction process is non-trivial: it requires intercepting the model's forward pass at specific points, capturing tensors, and saving them to disk in a format compatible with the training pipeline.

The assistant had already attempted one approach using vLLM (another serving framework) for hidden state extraction, but had pivoted to SGLang after discovering that vLLM's EAGLE-3 integration yielded only a ~15% acceptance rate—far too low to be useful. Now, with SGLang running and tuned to 90 tok/s single-stream performance, the next critical step was extracting hidden states from this new serving setup.

The Problem: How to Capture Hidden States from SGLang

The assistant faced a classic engineering dilemma. SGLang's model forward pass is tightly coupled to its internal ForwardBatch objects—the serving framework manages batching, scheduling, and tensor parallelism internally. Extracting hidden states meant either:

  1. Patching the server code to dump hidden states during normal inference (Approach C)
  2. Writing a standalone offline extraction script that loads the model outside the serving framework
  3. Finding an existing API that already supported this functionality The assistant had been leaning toward Approach C, planning to write a server-side patch that would intercept hidden states at layers [3, 31, 59] during prefill and save them as binary .pt files. This approach, while workable, carried significant risks: it required deep understanding of SGLang's internals, risked destabilizing the serving process, and would need careful synchronization to avoid data races.

The Discovery: SGLang Already Has This Feature

Message [msg 3283] begins with the assistant's realization:

SGLang already has a return_hidden_states API! When enable_return_hidden_states is set on the server AND return_hidden_states=True on the request, it captures FULL hidden states. And when capture_aux_hidden_states is True (EAGLE-3 mode), it concatenates the aux hidden states along the last dim and returns them.

This discovery came from tracing through SGLang's codebase. The assistant had been examining the model's forward pass, the logits processor, and the scheduler, looking for how hidden states flow through the system. In the process, it found two key mechanisms:

  1. enable_return_hidden_states: A server-level flag that activates the hidden states capture pipeline
  2. capture_aux_hidden_states: An EAGLE-3-specific mode that captures intermediate layer states The assistant immediately understood the implications:
This means I can: 1. Launch SGLang with --enable-return-hidden-states AND with layers_to_capture set (via EAGLE-3 config or monkey-patch) 2. Send requests via HTTP API with return_hidden_states=True 3. The hidden states come back in the response

This is a textbook example of the value of infrastructure investigation. Rather than building a custom solution from scratch, the assistant discovered that the framework already supported the exact functionality needed—it just wasn't obvious from the documentation or command-line flags.

Verifying the Mechanism

The assistant didn't stop at the discovery. It immediately verified how hidden states are returned in the HTTP response by examining the scheduler output processor code:

if req.return_hidden_states and logits_output.hidden_states is not None:
    req.hidden_states.append(
        logits_output.hidden_states[i].cpu().clone().tolist()
    )

This code snippet reveals the exact mechanism: when a request has return_hidden_states=True and the logits processor produces hidden states, they are converted to CPU tensors, cloned (to avoid reference issues), and converted to Python lists via .tolist(). These lists are then included in the HTTP response to the client.

The .tolist() conversion is notable—it means the hidden states are serialized as JSON-compatible nested lists, which could be very large for full hidden state vectors. For a model with 7168-dimensional hidden states across multiple layers, this could produce enormous JSON payloads. This is a potential performance concern that the assistant would need to address.

Assumptions and Their Overturning

This message reveals several important assumptions that were either confirmed or overturned:

Assumption 1: "SGLang doesn't have a built-in hidden states API." This was the implicit assumption driving the entire Approach C patching strategy. The discovery proved this assumption wrong.

Assumption 2: "We need to write custom code to capture hidden states." While partially true (some configuration and potentially a monkey-patch for layers_to_capture would still be needed), the heavy lifting was already done by SGLang.

Assumption 3: "The hidden states API, if it exists, would return raw tensors." The actual implementation returns Python lists via .tolist(), which has implications for data volume and serialization overhead.

Assumption 4: "EAGLE-3's capture_aux_hidden_states mode and the general return_hidden_states mode are separate systems." The assistant discovered they work together: when both are active, the aux hidden states (intermediate layers) are concatenated along the last dimension and returned through the same pipeline.

The Thinking Process Visible in the Message

The message reveals a clear, methodical thought process:

  1. Discovery: The assistant finds the return_hidden_states API through code exploration
  2. Synthesis: It immediately connects this discovery to the existing capture_aux_hidden_states mechanism for EAGLE-3
  3. Planning: It formulates a three-step plan (launch with flag, send requests, receive states)
  4. Verification: It doesn't just assume the API works—it traces through the code to confirm how hidden states are returned in the response This pattern of "discover → synthesize → plan → verify" is characteristic of effective debugging and engineering. The assistant doesn't celebrate the discovery and move on; it immediately validates the mechanism by reading the actual implementation code.

Input Knowledge Required

To understand and act on this message, the assistant needed:

  1. SGLang's codebase structure: Knowledge of where model definitions, logits processors, and schedulers live
  2. The EAGLE-3 training pipeline: Understanding that hidden states from specific layers are needed, and what format they must be in
  3. The existing extraction scripts: Awareness of 02_extract_hidden_states.py and 02b_extract_hidden_states_sglang.py and their approaches
  4. HTTP API patterns: Understanding that return_hidden_states=True would be passed as a request parameter
  5. Tensor serialization: Knowledge that .tolist() converts tensors to Python lists, and the implications for data size
  6. The server launch configuration: Knowing which flags are available (--enable-return-hidden-states) and how they interact

Output Knowledge Created

This message produced several valuable pieces of knowledge:

  1. A clear path forward: The assistant now knows exactly how to extract hidden states from SGLang without writing a custom patch
  2. Code confirmation: The scheduler output processor code confirms the exact mechanism and data format
  3. Risk reduction: The approach is non-invasive—it uses existing APIs rather than modifying server code
  4. A testable hypothesis: The assistant can now test whether --enable-return-hidden-states works with the current SGLang build and the Kimi-K2.5 model

The Broader Impact

This discovery had far-reaching consequences for the project. Instead of spending hours or days developing, testing, and debugging a server-side patch, the assistant could:

  1. Launch SGLang with --enable-return-hidden-states
  2. Write a simple client script that sends requests with return_hidden_states=True
  3. Collect the hidden states directly from the HTTP response The approach was also non-invasive—it didn't modify any SGLang source code, making it compatible with future SGLang updates and avoiding the maintenance burden of custom patches. In the subsequent messages (not shown here), the assistant would go on to implement this approach, successfully extracting 10K samples of hidden states from SGLang and using them to train a new EAGLE-3 drafter from scratch with dramatically better accuracy than the previous attempt.

Conclusion

Message [msg 3283] is a masterclass in the value of code exploration before building. Rather than charging ahead with a complex custom patch, the assistant took the time to trace through SGLang's existing code paths, discovering that the framework already supported the exact functionality needed. This saved what could have been days of engineering effort and produced a cleaner, more maintainable solution.

The message also demonstrates the importance of verification: the assistant didn't just note the existence of the API flags—it traced through the scheduler code to confirm exactly how hidden states are serialized and returned. This level of thoroughness is what separates robust engineering from fragile hacks.

For anyone working with large language model serving frameworks, the lesson is clear: before writing custom infrastructure, spend the time to understand what your framework already provides. The answer might be hiding in plain sight, buried in the codebase but ready to be discovered.