The Hidden State Extraction Problem: Engineering Precision in LLM Training Pipelines

Introduction

In the complex ecosystem of large language model (LLM) deployment and fine-tuning, few tasks are as technically demanding as extracting intermediate representations from a running inference server. Message 3294 captures a pivotal moment in such an endeavor: an AI assistant reasoning through the intricacies of capturing hidden states from a production SGLang server hosting the Kimi-K2.5 model, destined for use in training an EAGLE-3 speculative decoding drafter. This single message, dense with architectural insight and practical engineering judgment, reveals the layered decision-making required when adapting a serving system for a purpose it was not designed to serve.

The message sits at a critical juncture in a larger workflow. The team has already trained one EAGLE-3 drafter that failed to achieve meaningful speedup due to a ~15% acceptance rate. The root cause was traced to the training data: hidden states extracted via vLLM, which used a different model architecture version or quantization scheme than the SGLang server now running the production model. The solution was clear — re-extract hidden states using SGLang's own forward pass — but the implementation was anything but straightforward.

The Core Problem: Tensor Parallelism and Hidden State Integrity

The message opens with a crucial architectural verification:

The hidden_states tensor is full hidden_size=7168 on every TP rank (standard transformer TP approach — activations are replicated, weights are sharded). So dumping from any TP rank gives the correct values.

This statement reflects a deep understanding of how tensor parallelism (TP) works in transformer models. In standard TP implementations, the weights of linear layers are sharded across GPUs — each rank holds a slice of each weight matrix. However, the activations (hidden states) are typically replicated: each rank computes its portion of the output, then an all-reduce operation combines them into a complete tensor that every rank possesses. This is distinct from pipeline parallelism or sequence parallelism, where different ranks hold different parts of the data.

The assistant's verification here is critical. If hidden states were sharded across TP ranks (as they are in some model-parallel attention implementations), dumping from a single rank would produce only a fraction of each hidden state vector, corrupting the training data. By confirming that SGLang's DeepseekV2 implementation uses replicated activations, the assistant establishes that a simple single-rank dump strategy will work.

This assumption is tested implicitly by examining the codebase — the assistant has previously inspected the aux_hidden_states collection mechanism in the DeepseekV2 model forward pass (messages 3275-3278) and confirmed that hidden states are stored as full tensors before any TP reduction. The reasoning chain shows careful cross-referencing between the model implementation and the TP architecture.

The Batching Problem: Disentangling Requests in Flight

Having resolved the TP question, the assistant immediately identifies a more subtle challenge:

Now, there's another issue: in the serving loop, multiple requests may be batched, so the hidden_states tensor has shape [total_tokens_in_batch, 7168]. We need to know the boundaries of each request to separate them.

This is the crux of the difficulty. SGLang, like any production inference engine, optimizes throughput by batching multiple requests together. When requests share the same prefill phase (processing input tokens), their token sequences are concatenated into a single tensor: [total_tokens_in_batch, hidden_size]. The hidden states captured during this batched forward pass contain interleaved data from multiple requests, with no simple way to demultiplex them without access to the scheduler's per-request token counts.

The assistant's proposed solution is elegant in its simplicity:

The simplest solution: run with max_running_requests=1 to ensure only one request is in-flight at a time. Then the captured hidden states correspond to exactly that one request.

This is a classic engineering trade-off: sacrificing throughput (processing one request at a time instead of batching) for correctness (guaranteeing that captured hidden states belong to a single, identifiable request). For a data extraction pipeline processing thousands of samples, this is acceptable — the bottleneck is not server throughput but the quality and correctness of the training data.

However, the assistant immediately identifies a complication: even with max_running_requests=1, a mixed prefill+decode batch could occur. If a previous request is still in its decode phase while a new request begins prefill, the hidden state tensor would contain both the new request's prefill tokens and the old request's decode token. The assistant notes:

However, for a mixed prefill+decode batch, there could still be a decode token from a previous request in progress. With max_tokens=1 (generate only 1 token), the request finishes quickly after prefill. But we should ensure only one request processes at a time.

This reveals an important nuance: the max_running_requests=1 setting limits how many requests are accepted into the scheduler, but it does not guarantee that the scheduler won't batch a prefill from one request with a decode from another if the timing overlaps. The assistant's mitigation is to set max_tokens=1, ensuring each request generates exactly one token and completes almost immediately after its prefill pass. This minimizes the window for cross-request contamination.

The CaptureHiddenMode Problem: Prefill vs. Decode

The assistant then identifies yet another subtlety in SGLang's hidden state capture mechanism:

The issue is that CaptureHiddenMode.FULL captures hidden states for ALL forward passes (prefill + decode). But for a single request with max_tokens=1, the first forward is prefill (all tokens) and the second is decode (1 token). We want the prefill one.

SGLang's CaptureHiddenMode system, which the assistant has been investigating in previous messages (3282-3284), is designed to return hidden states through the HTTP API. When set to FULL, it captures hidden states at every forward pass — both the prefill pass (which processes all input tokens at once) and each subsequent decode pass (which processes one new token at a time). For EAGLE-3 training, only the prefill hidden states are needed, since they represent the complete input sequence's representations across all layers.

The assistant's solution is to modify the forward pass to save hidden states only when the forward mode is prefill:

Let me just write the patch — I'll save the hidden states from the prefill pass only (when the forward mode is prefill).

This requires understanding SGLang's ForwardMode enum and how it distinguishes between prefill (called "extend" in SGLang's terminology) and decode phases. The assistant immediately begins investigating this:

[bash] ssh root@10.1.230.174 "grep -n 'forward_mode\|ForwardMode\|is_prefill\|is_extend\|is_decode' /root/sglang/python/sglang/srt/managers/schedule_batch.py | head -20"

This command searches the SGLang scheduler code for the relevant mode definitions, showing the assistant's methodical approach to understanding the codebase before writing a patch.

The Broader Strategy: Approach C in Context

This message represents the culmination of a multi-step investigation into hidden state extraction strategies. Earlier messages (3279-3285) explored three approaches:

Assumptions and Their Validity

The message rests on several key assumptions:

  1. Hidden states are replicated across TP ranks. This is verified through code inspection but not empirically tested. If SGLang's implementation differs from the standard transformer TP pattern for this specific model architecture, the assumption could fail silently, producing corrupted training data.
  2. max_running_requests=1 with max_tokens=1 is sufficient to prevent batching contamination. This assumes the scheduler respects these limits strictly and that no edge cases (e.g., internal scheduler prefetching, background tasks) introduce additional tokens into the batch.
  3. The prefill forward pass captures all input positions. This is correct for causal attention with left-padding or no padding, but could be wrong if the model uses right-padding or if the scheduler truncates sequences.
  4. Dumping from a single TP rank is sufficient. While architecturally correct, this assumes no silent data corruption or numerical divergence across ranks (unlikely but possible in practice).

Knowledge Required and Created

To fully understand this message, a reader needs:

Conclusion

Message 3294 exemplifies the kind of deep, systems-level thinking required when adapting production ML infrastructure for research purposes. It is not a message about writing code — it is a message about understanding a complex system well enough to modify it safely. The assistant's reasoning moves from architectural invariants (TP replication) through operational constraints (request batching) to implementation details (ForwardMode enum), building a complete mental model of the extraction problem before writing a single line of the patch.

This message also illustrates a broader truth about LLM engineering: the hardest problems are often not about model architecture or training algorithms, but about the data plumbing between them. The difference between a working EAGLE-3 drafter and a broken one was traced to a mismatch in hidden state sources (vLLM vs. SGLang). Getting the extraction right required understanding not just what hidden states to capture, but when and how the serving engine computes them — knowledge that lives at the intersection of systems engineering and deep learning research.