The Radix Cache Problem: A Critical Insight in Hidden State Extraction for EAGLE-3 Training

In the high-stakes world of speculative decoding for large language models, every detail matters. A single incorrect assumption about how the inference engine processes tokens can silently corrupt an entire training dataset, wasting days of compute and rendering a drafter model useless. Message [msg 3381] captures the precise moment when an engineer discovered such a subtle pitfall — and made the critical decision to address it before it could poison the EAGLE-3 training pipeline.

The Context: Building a Better Drafter

The broader session (Segment 25) was focused on a sophisticated engineering goal: training a new EAGLE-3 speculative decoding drafter for the Kimi-K2.5 model, a large Mixture-of-Experts architecture running on 8 NVIDIA RTX PRO 6000 Blackwell GPUs (SM120 architecture). The assistant had already achieved impressive single-stream throughput of 90 tok/s on SGLang, surpassing vLLM's 82.5 tok/s, and had developed a non-invasive server-side patch to capture intermediate hidden states during model inference.

The core idea behind EAGLE-3 is elegant: train a lightweight "drafter" model that predicts not just the next token, but the hidden state at the next position, allowing the large model to verify multiple candidate tokens in parallel. For this to work, the training data must consist of ground-truth hidden states — the actual intermediate representations produced by the base model at specific layers. If these hidden states are wrong, the drafter learns the wrong thing.

The assistant had written a custom patch for SGLang's kimi_k25.py model file that intercepted the forward pass during prefill (the EXTEND phase) and saved the hidden states at layers 3, 31, and 59 to disk as PyTorch tensor files. After successfully launching the patched server and verifying it responded to health checks, the moment of truth arrived: did the hidden state extraction actually work?

The Moment of Discovery

Message [msg 3381] opens with a moment of vindication:

13 tokens in, 13 tokens dumped — perfect match! The previous 1-token result was because of radix cache hitting on a previously sent prompt. With 0 cached tokens, we get the full hidden states for all 13 input positions.

This conclusion was reached after a careful diagnostic chain. In the preceding message ([msg 3380]), the assistant had sent a 13-token prompt and received a dump showing exactly 13 hidden states across all four saved tensors (three auxiliary layers plus the final layer). The shapes confirmed it: torch.Size([13, 7168]) for each tensor, where 7168 is the hidden dimension of Kimi-K2.5.

But this success was preceded by a confusing failure. Earlier, a test with 10 prompt tokens had produced a dump of only 1 token. The assistant's reasoning in [msg 3381] reconstructs what happened: the radix cache — SGLang's mechanism for reusing KV cache entries across requests with common prefixes — had cached the prompt from a previous request. When the new request arrived, SGLang only needed to process the uncached portion, and the hidden state capture code only recorded those new tokens.

Understanding the Radix Cache Problem

The radix cache is a performance optimization in SGLang that stores KV cache entries in a tree structure, allowing requests with shared prefixes to reuse computation. This is normally a huge win for serving workloads where many requests share system prompts or conversation history. But for hidden state extraction, it's a disaster.

Here's why: the assistant's patch captures hidden states during the EXTEND forward pass — the phase where new tokens are processed against the existing KV cache. When the radix cache is enabled, SGLang skips the forward pass for any tokens already in the cache. The hidden states for those cached tokens are never materialized in the forward pass, so they're never dumped. The result is a partial dump containing only the non-cached suffix of the prompt.

For EAGLE-3 training, this is unacceptable. The drafter needs to learn to predict hidden states at every position in the sequence. If the training data only contains hidden states for the last few tokens of each prompt, the drafter will never learn to handle early positions, and its acceptance rate during speculative decoding will suffer accordingly.

The assistant's reasoning in [msg 3381] shows a clear understanding of this:

This is crucial: we need to make sure there's no caching between requests. For extraction, each request should have unique input so the radix cache doesn't interfere. Since our tokenized training data consists of different conversations, this should naturally be the case.

This is a reasonable first thought: if every prompt in the training dataset is a unique conversation, then no two requests share a prefix, and the radix cache will never match. But then the assistant catches a subtler edge case:

But wait — there's also the issue that if the first few tokens of a sequence match a previous sequence (e.g., system prompt), those tokens would be cached and the EXTEND forward would only process the remaining tokens. For EAGLE-3 training, we need hidden states for ALL positions.

This is the critical insight. Even if the training conversations are unique, they might share a common system prompt prefix (e.g., "You are a helpful assistant..." or a role-play preamble). If the radix cache matches on this prefix, the first N tokens of every sequence would be silently dropped from the hidden state dump. The drafter would never learn to predict hidden states for the system prompt region, which could be important for maintaining conversation coherence.

The Decision: Disable Radix Cache

The assistant's conclusion is decisive:

The safest approach: disable radix cache. Let me check if there's a flag for that:

This decision represents a tradeoff. Disabling radix cache will make extraction slower — each request must process all tokens from scratch, increasing latency and GPU memory bandwidth usage. But correctness trumps performance here. A corrupted training dataset would produce a drafter with degraded acceptance rates, potentially wasting the entire training run.

The assistant immediately checks the server launch arguments to confirm the current state:

[assistant] ssh root@10.1.230.174 "grep 'disable_radix_cache' /data/eagle3/synth_10k/sglang_hs_dump_v2.log | head -1"
[2026-02-23 10:23:16] server_args=ServerArgs(model_path='/shared/kimi-k2.5-int4', tokenizer_path='/shared/kimi-k2.5-int4', tokenizer_mode='auto', ...)

The log line confirms that --disable-radix-cache was not set, meaning radix cache was active. The next step (visible in the following message [msg 3382]) would be to restart the server with the flag enabled.

Assumptions, Knowledge, and Reasoning

This message reveals several layers of the assistant's thinking process:

Input knowledge required: To understand this message, one needs to know what the radix cache is and how it works in SGLang, understand the EXTEND vs DECODE forward pass distinction, know the EAGLE-3 training data format (hidden states at specific layers for every token position), and be familiar with the Kimi-K2.5 model architecture (hidden size 7168, MLA attention).

Assumptions made: The assistant initially assumed that unique conversation data would naturally avoid radix cache interference. This was a reasonable assumption that proved only partially correct — it would work for the training data itself but could fail for edge cases like shared system prompts. The assistant also assumed that the server's warmup requests (which produced req_0, req_1, req_2) wouldn't interfere with the extraction counter, which required additional handling.

Output knowledge created: This message produced the verified knowledge that the hidden state extraction patch works correctly when radix cache is not interfering, and that radix cache must be disabled for correct extraction. It also established the diagnostic procedure: send a unique prompt, check cached_tokens in the response metadata, and verify that the dump tensor shapes match the prompt token count.

Mistakes and corrections: The earlier test with prompt_token_ids (which returned 400 errors) and the initial 1-token dump result were both corrected through systematic debugging. The assistant methodically ruled out tokenizer issues, API endpoint differences, and finally identified radix cache as the root cause.

The Broader Significance

This message exemplifies a class of problems that arise when repurposing inference infrastructure for training data generation. SGLang's radix cache is designed for production serving, where throughput and latency are paramount. Using the same server for hidden state extraction requires understanding which optimizations are safe and which silently corrupt the data.

The decision to disable radix cache also has downstream implications. With radix cache disabled, the extraction of 10,000 samples (the target mentioned in the segment summary) would be slower and consume more GPU memory bandwidth. But the alternative — training a drafter on partial hidden states — would likely produce a model with ~25% acceptance rate (like the previous broken drafter) rather than the dramatically better accuracy the team was aiming for.

In the end, this message captures the moment when a subtle systems-level insight saved the entire training pipeline from a silent, hard-to-diagnose failure mode. It's a reminder that in ML engineering, understanding the inference engine's internals is just as important as understanding the model architecture itself.