Designing a Server-Side Hidden State Extraction Patch for SGLang

In message [msg 3298] of a lengthy coding session focused on building an EAGLE-3 speculative decoding system for the Kimi-K2.5 large language model, the assistant reaches a critical design decision: how to extract intermediate hidden states from SGLang's serving infrastructure to generate training data for a new EAGLE-3 drafter. This message represents a pivotal moment where architectural understanding, practical constraints, and system-level design converge into a concrete implementation plan.

The Problem: Extracting Hidden States at Scale

The broader session context reveals a multi-week effort to train an EAGLE-3 speculative decoding drafter for Kimi-K2.5, a massive Mixture-of-Experts model running on 8 NVIDIA RTX PRO 6000 Blackwell GPUs. The assistant had previously extracted hidden states using vLLM, but those proved inadequate — the trained drafter achieved only ~25% acceptance rate, far below what was needed for useful speedup. A pivot to SGLang had been made, and now the assistant needed to extract 10,000 high-quality hidden state samples from SGLang to retrain the drafter from scratch.

The core challenge is architectural: SGLang is a high-performance serving engine designed for throughput, not for introspecting model internals. Hidden states — the intermediate activations at specific transformer layers — are essential for training EAGLE-3's lightweight draft model, which learns to predict likely next tokens given these internal representations. But extracting them from a running server without disrupting its operation or corrupting the data requires careful engineering.

The Reasoning Journey: Three Approaches Considered

The assistant's thinking, visible across the preceding messages ([msg 3278] through [msg 3297]), reveals a systematic exploration of three approaches before arriving at the final design.

Approach A: Use SGLang's built-in return_hidden_states API. The assistant discovered that SGLang already has an enable_return_hidden_states server flag and a return_hidden_states=True request parameter. However, deeper investigation revealed a fatal flaw: the server converts hidden states to Python lists via .tolist() and returns them as JSON in the HTTP response. For a 2048-token sequence with hidden_size=7168 across 3 layers, that's approximately 176 MB of float data per sample, serialized into verbose JSON text. Scaling this to 10,000 samples would be prohibitively slow and memory-intensive.

Approach B: Write a standalone offline extraction script. The assistant considered loading the model directly (using SGLang's weight loading) and running a forward pass without the serving infrastructure. This avoids the serialization problem entirely. But it runs into the "ForwardBatch complexity issue" — SGLang's model forward pass is tightly coupled to its batching infrastructure, and replicating that outside the server would require either significant code duplication or deep understanding of SGLang's internals.

Approach C: Patch the running server to dump hidden states to disk as binary files. This became the chosen approach. The idea is elegant: modify the model's forward method to write hidden states directly to /dev/shm/ as PyTorch .pt files during prefill, bypassing the JSON serialization entirely. The HTTP API is used only as a trigger — the client sends a request, the server processes it, dumps the hidden states to shared memory, and the client reads the binary files after the response returns.

The Design Decisions in Message 3298

The message itself crystallizes these decisions into a concrete implementation. The assistant writes:

"Let me write the complete approach: 1. Patch file — monkey-patches DeepseekV2Model to dump hidden states on prefill (EXTEND mode) 2. Extraction script — sends requests one at a time via HTTP, reads dumped files"

Several key design decisions are embedded here:

Decision 1: Monkey-patching over direct modification. Rather than modifying the SGLang source code directly (which would require rebuilding the container or managing git diffs), the assistant opts for a monkey-patch file. This is a pragmatic choice for a research environment where the server is already running and stability is paramount. The patch can be applied at runtime without restarting the server with modified code.

Decision 2: Dumping only during prefill (EXTEND mode). The assistant correctly identifies that EAGLE-3 training only needs hidden states from the prefill pass — the forward pass that processes all input tokens simultaneously. The decode pass (generating one token at a time) produces hidden states too, but those are for individual generated tokens, not the full input sequence. By filtering on forward_mode == EXTEND, the patch avoids dumping redundant decode-time states.

Decision 3: Using /dev/shm/ as the dump location. This is a Linux shared memory filesystem, providing fast, RAM-based storage that avoids disk I/O latency. For tensors totaling potentially hundreds of gigabytes across 10,000 samples, this is critical for performance.

Decision 4: Single-request batching. The assistant notes the complication that "multiple requests may be batched, so the hidden_states tensor has shape [total_tokens_in_batch, 7168]." The solution is to run with max_running_requests=1 to ensure only one request is in-flight at a time, making the hidden states correspond unambiguously to that single request.

Assumptions and Their Implications

The message rests on several assumptions that deserve scrutiny:

Assumption 1: Hidden states are identical across all TP ranks. The assistant verifies that with tensor parallelism (TP=8), the hidden_states tensor is the full hidden_size=7168 on every rank. This is correct for standard transformer TP where activations are replicated rather than sharded. The verification is done by reasoning about the code structure rather than empirical testing.

Assumption 2: The counter file mechanism prevents race conditions. The patch design includes "a counter file so the client knows when the dump is ready." This assumes that file system operations on /dev/shm/ are atomic enough for this signaling mechanism to work reliably under concurrent access.

Assumption 3: max_running_requests=1 is sufficient to prevent batching. While this prevents multiple requests from being processed simultaneously, SGLang's scheduler might still batch a prefill with a decode from a previous request. The assistant acknowledges this but notes that with max_tokens=1, the request finishes quickly after prefill, minimizing the window for mixed batches.

Assumption 4: The hidden states captured by SGLang's CaptureHiddenMode.FULL during prefill correspond exactly to the EAGLE-3 training format. The assistant had previously verified ([msg 3288]) that the vLLM-extracted format uses layers [2, 30, 58, 60] with shapes [seq_len, 7168]. The SGLang patch targets layers [3, 31, 59] (which correspond to the same values — the "output of layer i" equals the "input to layer i+1"). This mapping is correct but subtle, and any off-by-one error would produce incorrect training data.

Input Knowledge Required

To fully understand this message, one needs:

  1. SGLang architecture knowledge: Understanding of the DeepseekV2Model class, the forward method signature, ForwardMode enum, tensor parallelism, and the CaptureHiddenMode mechanism.
  2. EAGLE-3 training pipeline knowledge: Familiarity with the hidden state format (input_ids, hidden_states list of 4 tensors, loss_mask), the layer indices used, and the bfloat16 dtype requirement.
  3. System-level Linux knowledge: Understanding of /dev/shm/, file-based signaling between processes, and the performance characteristics of binary vs. JSON serialization.
  4. PyTorch distributed concepts: Knowledge of tensor parallelism and how hidden states are (or aren't) sharded across TP ranks.

Output Knowledge Created

This message produces several artifacts of knowledge:

  1. A design document for the extraction pipeline: The patch approach, the extraction script architecture, and the signaling mechanism are all specified here.
  2. A concrete implementation plan: The assistant writes a patch file (sglang_hs_dump_patch.py) that will be applied to the running SGLang server.
  3. Validation criteria: The approach implicitly defines what "success" looks like — binary .pt files appearing in /dev/shm/sglang_hs/ after each API request, with the correct tensor shapes and dtypes.
  4. A reusable pattern: The monkey-patching approach for extracting model internals from SGLang is a general technique that could be applied to other models or other internal states beyond hidden states.

The Thinking Process Visible in the Message

The assistant's reasoning in this message is notably systematic. It begins with a summary of what was learned ("Good. Now I have everything I need"), then structures the approach into two clear components (patch file and extraction script), enumerates the patch's behavior in bullet points, and finally writes the actual code.

What's particularly interesting is the parallel between the assistant's reasoning and software engineering best practices. The assistant:

Conclusion

Message [msg 3298] represents a critical design inflection point in a complex ML engineering effort. The assistant synthesizes knowledge from codebase exploration, prior extraction attempts, and system-level constraints into a clean, practical solution. The decision to patch the running server rather than use the built-in API or build a standalone extraction pipeline reflects a mature understanding of the tradeoffs between development speed, runtime performance, and data quality. The resulting extraction pipeline would go on to successfully produce 10,000 samples of hidden states (924 GB) with zero errors, enabling the training of a dramatically better EAGLE-3 drafter.