The Hidden State Patch: Instrumenting SGLang's Forward Pass for EAGLE-3 Training

In the middle of a complex multi-day effort to deploy speculative decoding for the Kimi-K2.5 large language model, a single bash command represents the culmination of extensive reasoning about model internals, server architecture, and the delicate boundary between serving and training infrastructure. Message [msg 3323] is deceptively simple:

[assistant] [bash] ssh root@10.1.230.174 "python3 /root/eagle3-train/apply_hs_dump_patch.py /root/sglang/python/sglang/srt/models/deepseek_v2.py"
Successfully patched /root/sglang/python/sglang/srt/models/deepseek_v2.py
  - Added hidden state dump imports
  - Added _hs_dump init in DeepseekV2Model.__init__
  - Added dump logic in DeepseekV2Model.forward
  - Auto-enabled capture_aux_hidden_states in CausalLM

To use: set SGLANG_HS_DUMP_DIR=/dev/shm/sglang_hs before starting server

This message executes a Python script that surgically modifies the SGLang serving engine's model implementation to capture intermediate hidden states during inference. It is the pivot point where the assistant transitions from using the model for serving to instrumenting it for training data generation. Understanding why this particular approach was chosen, and the reasoning that led to it, reveals the deep entanglement between inference engines and training pipelines in modern large language model development.

The Motivation: Why Hidden States Matter

The broader context is the training of an EAGLE-3 speculative decoding drafter for Kimi-K2.5, a 1-trillion-parameter mixture-of-experts model. EAGLE-3 works by training a lightweight "drafter" model that predicts multiple future tokens in parallel, conditioned on the hidden states of the base model. The drafter learns to predict what hidden states would look like after processing additional tokens, allowing the inference engine to speculate several tokens ahead and then verify them in a single forward pass.

To train this drafter, one needs ground-truth data: actual hidden states from the base model at specific layers, captured during real inference. The assistant had previously extracted 10,000 samples using vLLM, but the resulting drafter achieved only a ~15% acceptance rate — barely better than random. The pivot to SGLang (which loaded the model faster and achieved better throughput) meant that a new extraction pipeline was needed, this time using SGLang's model execution path.

The core challenge is that SGLang, like most inference engines, is optimized for production serving, not training data extraction. Its forward pass is a carefully optimized pipeline involving CUDA graphs, radix caching, and batched attention — none of which are designed to expose intermediate activations to Python code. The assistant needed to find a way to capture those activations without breaking the serving stack.

The Design Space: Three Approaches Considered

The assistant's reasoning, visible across messages [msg 3298] through [msg 3320], explores three fundamentally different approaches to hidden state extraction.

Approach A: Monkey-patching from outside. The initial idea was to write a Python script that monkey-patches the DeepseekV2Model.forward method at runtime, after the model is loaded but before serving begins. This would be clean and non-invasive — no source files modified. However, the assistant quickly realized a critical problem: the monkey-patch needs to be applied after the model object is instantiated but before any requests are served. This timing is difficult to achieve with SGLang's server startup sequence, which loads the model and immediately begins accepting requests.

Approach B: Offline extraction. The assistant considered bypassing the server entirely and writing a standalone script that loads the model weights using SGLang's infrastructure but runs a simple forward pass without the serving stack. This would avoid all the complexity of CUDA graphs and request scheduling. However, investigation revealed that the ForwardBatch object — which carries all the metadata needed for attention, KV cache management, and MoE routing — is deeply embedded in the model's forward pass (see [msg 3300]). Creating a fake ForwardBatch that satisfies all the model's internal requirements was deemed too complex and fragile.

Approach C: Direct source patching. The chosen approach: directly edit the deepseek_v2.py source file on the server to add hidden state dumping logic. This is invasive — it modifies the serving engine's code — but it guarantees that the dump logic runs in the exact same execution context as normal inference, with all the correct tensor shapes, attention masks, and cache states. The patch is controlled by an environment variable (SGLANG_HS_DUMP_DIR), so it can be enabled and disabled without code changes.

The Architecture of the Patch

The patch script (apply_hs_dump_patch.py) makes four targeted modifications to deepseek_v2.py:

  1. Add imports for os, torch, and pathlib at the top of the file, which are needed for file I/O and directory management within the forward pass.
  2. Initialize dump state in DeepseekV2Model.__init__: Adds self._hs_dump_dir (read from the environment variable) and a counter attribute. This runs once at model load time and sets up the dump directory.
  3. Insert dump logic in DeepseekV2Model.forward: This is the core of the patch. After the model computes the final hidden states (the output of the transformer stack, after the final norm), the patch checks if dump mode is enabled and if the current forward mode is EXTEND (prefill). If so, it saves the intermediate hidden states — captured at layers specified by layers_to_capture — as binary .pt files to the dump directory. Critically, this logic runs inside the forward pass, meaning it has access to the exact tensors flowing through the model.
  4. Auto-enable capture_aux_hidden_states in DeepseekV2ForCausalLM.__init__: The CausalLM wrapper class has a flag capture_aux_hidden_states that controls whether the model returns auxiliary hidden states alongside the final logits. Normally this flag is set to False and only enabled when EAGLE-3 mode is active (with a draft model configured). The patch forces it to True when the dump environment variable is set, ensuring the hidden states are computed and returned even without a draft model.

The KimiK25 Compatibility Problem

A subtle but critical issue emerged during the patch design: the Kimi-K2.5 model uses a custom wrapper class (KimiK25ForConditionalGeneration) that delegates to DeepseekV2ForCausalLM via self.language_model. The assistant checked whether this wrapper properly handles capture_aux_hidden_states and aux_hidden_states ([msg 3309]–[msg 3316]). The investigation revealed that kimi_k25.py does not reference these attributes at all — it simply calls self.language_model.forward() and returns the logits.

This meant the assistant had to verify that the hidden state flow works correctly through the delegation chain: KimiK25ForConditionalGeneration.forward → general_mm_embed_routine → DeepseekV2ForCausalLM.forward → DeepseekV2Model.forward. The dump happens at the innermost level (DeepseekV2Model.forward), which is called during the CausalLM's forward pass. The capture_aux_hidden_states flag on the CausalLM ensures that the hidden states are computed and stored (via logits_processor) even though the outer KimiK25 wrapper doesn't explicitly request them.

The CUDA Graph Problem

Another critical consideration was CUDA graphs. SGLang uses CUDA graph capture to accelerate inference: the entire forward pass is traced once and then replayed for every request, bypassing Python execution entirely. If CUDA graphs are enabled, the dump logic (which is pure Python code) would only execute during the initial graph capture, not during actual inference. The assistant correctly identified this and planned to launch the extraction server with --disable-cuda-graph (see [msg 3320]). This trades inference speed for correctness — without CUDA graphs, the server is slower, but every forward pass executes the Python dump logic.

Assumptions and Potential Pitfalls

The patch makes several assumptions that could fail in edge cases:

The Significance of This Message

Message [msg 3323] is a turning point in the session. Before it, the assistant had been benchmarking SGLang's serving performance, achieving 90 tok/s single-stream throughput — a 9% improvement over vLLM. After it, the assistant would restart the server with the patch enabled, extract 10K samples of hidden states, and begin training a new EAGLE-3 drafter from scratch.

The patch itself is a temporary, surgical modification — the assistant later kills the extraction server and restarts without the patch for production serving. But in that window, it enables a data pipeline that would prove dramatically more successful than the previous attempt: the new drafter achieved ~74% step 0 accuracy, ~64% step 1 conditional accuracy, and ~56% step 2 conditional accuracy, compared to the previous drafter's ~25% acceptance rate.

This message exemplifies a recurring pattern in ML engineering: the boundary between serving and training is porous. To improve inference performance through speculative decoding, one must reach back into the serving engine and extract the very data that training requires. The patch is a bridge between these two worlds — a temporary modification that turns a production inference server into a data generation tool, enabling the next iteration of the model to serve faster.