From 90 Tok/s to EAGLE-3: Building a Complete Speculative Decoding Pipeline for Kimi-K2.5

Introduction

The gap between a working inference server and a trained speculative decoding drafter is rarely a straight line. In Segment 25 of this opencode session, an AI assistant navigated a complex, multi-stage pipeline that touched nearly every layer of modern machine learning infrastructure: performance tuning of an inference engine, developing a non-invasive server-side patch for hidden state extraction, orchestrating a multi-terabyte data extraction pipeline, training a 2.6-billion-parameter draft model from scratch, debugging silent logging failures, and finally visualizing training metrics to confirm success. This article synthesizes the work across this segment, tracing the narrative arc from SGLang optimization through to the first validated training metrics of a new EAGLE-3 drafter for the Kimi-K2.5 language model — a 1-trillion-parameter Mixture-of-Experts (MoE) model deployed across eight NVIDIA RTX PRO 6000 Blackwell GPUs (SM120 architecture).

The Performance Foundation: Surpassing vLLM

The segment opened with a focused performance optimization effort. The assistant had previously deployed Kimi-K2.5 on both vLLM and SGLang, but the initial SGLang benchmark was disappointing: only 63.6 tok/s single-stream throughput, far below vLLM's 82.5 tok/s baseline. Since the EAGLE-3 pipeline required SGLang (vLLM's EAGLE-3 integration with MLA had proven broken in earlier segments), closing this performance gap was essential before any further work could proceed.

The assistant's investigation traced the root cause to NCCL communication overhead. In a tensor-parallel setup with 8 GPUs, every transformer layer requires all-reduce operations to synchronize activations across ranks. SGLang's default NCCL configuration was not optimized for the Blackwell GPU architecture and the specific NVLink topology.

The solution came through two key tuning levers. First, the assistant applied NCCL environment variables — NCCL_PROTO=LL (Low Latency protocol) and NCCL_ALGO=Ring (Ring all-reduce algorithm) — to optimize inter-GPU communication for single-stream workloads. The LL protocol prioritizes latency over bandwidth, which is ideal for the small message sizes typical of single-token decode. The Ring algorithm provides good bandwidth utilization across 8 GPUs. Second, the assistant added --num-continuous-decode-steps 4, a SGLang flag that batches multiple decode steps together, reducing scheduling and attention recomputation overhead.

The result, captured in [msg 3319], was a decisive 90.0 tok/s single-stream throughput — a 42% improvement over the untuned baseline and a 9% improvement over vLLM's 82.5 tok/s. The assistant celebrated the milestone but immediately noted the trade-off: at high concurrency (128 concurrent requests), throughput dropped slightly from 2,370 tok/s to 2,320 tok/s due to continuous decode step overhead. This was an acceptable cost given that the primary use case — hidden state extraction — is inherently single-stream.

Notably, the successful configuration used triton attention (the default for DeepSeek models on SM120) rather than FlashInfer, which had caused the server to hang on the Blackwell architecture. This was a forced choice by hardware compatibility, but it turned out to be the right one.

The Hidden State Problem: Why Extraction Was Necessary

With SGLang performing at peak, the assistant turned to the core challenge: extracting hidden states for EAGLE-3 training. The EAGLE-3 speculative decoding architecture requires a lightweight "drafter" model that predicts multiple future tokens in parallel, conditioned on intermediate hidden states from the base model. For Kimi-K2.5, the drafter needs hidden states from specific transformer layers — layers 3, 31, and 59 (corresponding to EAGLE-3's feature layers [2, 30, 58] plus a one-index offset) — captured during the prefill phase.

The assistant had previously extracted 10K samples using vLLM, but the resulting drafter achieved only a ~15% acceptance rate, making speculative decoding slower than running the base model directly (0.66× throughput). The pivot to SGLang was motivated by hope that a different serving stack would yield better hidden states — but SGLang had no built-in mechanism for large-scale hidden state extraction.

The assistant systematically explored three approaches ([msg 3284]–[msg 3301]). Approach A — using SGLang's existing return_hidden_states API — was rejected because it converts hidden states to Python lists via .tolist() and serializes them over HTTP as JSON, which would be catastrophically slow for thousands of samples. Approach B — offline extraction by loading the model weights outside the serving stack — was rejected after examining the ForwardBatch complexity; the model's forward pass is deeply entangled with KV cache management, attention backends, and tensor parallelism. Approach C — server-side patching of the running SGLang server to dump hidden states to disk as binary .pt files — was selected as the winning strategy.

The Patch: From v1 Hang to v2 Non-Invasive Success

The assistant's first attempt at the server-side patch (v1) modified SGLang's existing capture_aux_hidden_states mechanism. This approach seemed natural — SGLang already had infrastructure for capturing auxiliary hidden states during model forward passes. The patch set capture_aux_hidden_states = True and configured layers_to_capture = [3, 31, 59].

But the v1 patch proved disastrous. The server would load weights successfully (the log showed 217 lines of weight loading output) and then go silent. The HTTP port never opened. The scheduler processes were running but stuck in busy loops, consuming high CPU. The health endpoint returned "Connection refused."

The assistant spent several messages diagnosing this hang ([msg 3335] through [msg 3350]), checking process file descriptors, examining strace output, and comparing log line counts against a known-working server configuration. The key insight came when the assistant realized that enabling capture_aux_hidden_states changed how the CausalLM model class interpreted the return value from the DeepseekV2Model.forward method. Specifically, the line hidden_states, aux_hidden_states = hidden_states would unconditionally try to unpack the model's return value as a tuple. But during certain forward modes — particularly warmup and idle passes — the model might return only a single tensor, causing a silent failure or deadlock.

Rather than continuing to debug the interaction between capture_aux_hidden_states and SGLang's warmup logic — a path that could have led deep into the framework's internals — the assistant made a strategic decision to pivot. It reverted the v1 patch entirely and designed a fundamentally different approach.

The v2 patch ([msg 3357], [msg 3359]) takes a "non-invasive" strategy. Instead of modifying SGLang's existing hidden state capture mechanism, it adds independent capture logic directly within the DeepseekV2Model.forward loop. The patch:

  1. Reads an environment variable (SGLANG_HS_DUMP_DIR) to determine where to save hidden states
  2. Checks, at each layer in the forward loop, whether the current layer index is in the capture set {3, 31, 59}
  3. If so, computes hidden_states + residual (the post-attention residual stream) and saves it to a buffer
  4. After the forward pass completes, dumps the captured states to disk as binary .pt files Crucially, this approach does not touch capture_aux_hidden_states, layers_to_capture, or any other SGLang configuration. The server operates completely normally. The hidden state capture is a side effect that runs in parallel with the normal forward computation, not a modification of the server's control flow. The patch also required careful handling of CUDA graphs. When CUDA graphs are enabled, the forward pass is captured as a compiled graph and replayed — any Python code (like the dump logic) only executes during the initial capture, not during replay. The assistant explicitly launched the server with --disable-cuda-graph to ensure the dump logic ran on every forward pass ([msg 3329]).

Verification: The Moment of Proof

After applying the v2 patch and launching the server, the assistant needed to confirm it was actually working. The verification came in two stages.

First, in [msg 3333], the assistant grepped the server log for "HS_DUMP" lines and found:

[HS_DUMP] Model init: dump dir=/dev/shm/sglang_hs, layers_to_capture=[3, 31, 59]
[HS_DUMP] CausalLM: capture_aux_hidden_states=True

These two log lines confirmed that the initialization hooks fired correctly — the environment variable was read, the dump directory was created, the layer configuration was set, and the upstream flag was enabled. This was the gate that separated a failed pipeline from a successful one.

Second, in [msg 3365], the assistant confirmed the server was fully operational. Port 8000 was listening, the health endpoint returned empty (SGLang's convention for "healthy"), and the get_model_info endpoint returned a complete JSON response confirming the Kimi-K2.5 model was loaded with the correct architecture. This was the moment the pipeline unblocked.

The final verification came in [msg 3367], where the assistant sent a test request and inspected the dump output. The tensor shapes confirmed correct extraction: each .pt file had shape [seq_len, 7168] with dtype bfloat16 — exactly matching the model's hidden size and precision. The three auxiliary files (aux_0.pt, aux_1.pt, aux_2.pt) corresponded to the three captured layers, and final.pt contained the final hidden state.

The Radix Cache Problem: A Critical Insight

One of the most consequential discoveries in this segment was the radix cache problem. SGLang uses a radix cache to accelerate inference by caching KV cache entries for common prefixes across requests. When a new request shares a prefix with a previous request, the cached tokens are reused, and only the new tokens go through the forward pass. This is excellent for inference throughput but catastrophic for hidden state extraction: if a prefix is cached, the hidden states for those prefix tokens are never computed, and thus never captured by the extraction patch.

The assistant identified this issue when a test request with 13 tokens produced only 12 hidden state vectors — the first token had been cached from a previous request. The realization was immediate and decisive: the radix cache must be disabled for correct extraction. The server was restarted with --disable-radix-cache and --disable-cuda-graph, ensuring that every token in every request would pass through the full forward computation.

This decision had operational consequences. Restarting the SGLang server on an 8-GPU machine with a 200B+ parameter model took approximately 9 minutes — a significant idle period. The assistant used this time productively, gathering statistics on the training dataset (10,000 samples, 21 million tokens, estimated 1.2 TB of hidden states) and checking available disk space (1.9 TB free on /data). This opportunistic multitasking — using unavoidable latency windows for preparatory work — is a hallmark of efficient engineering workflow.

The 10K-Sample Extraction: Orchestrating at Scale

With the server correctly configured, the assistant launched the full 10K-sample extraction. The extraction script (02b_extract_hidden_states_sglang.py) sent each tokenized prompt to the SGLang server, which ran the patched forward pass, dumped intermediate hidden states to /dev/shm/, and the script collected and saved them to persistent storage.

The extraction was not without its own debugging challenges. The first launch attempt failed because the output directory /data/eagle3/synth_10k_sglang/ did not exist — the shell redirection for the log file failed before the Python script ever started. A single mkdir -p command unblocked the entire pipeline, a reminder that in complex ML systems, the most mundane infrastructure detail can be the difference between a running pipeline and a stalled one.

A more subtle issue emerged when the assistant checked the extraction log and found it completely empty — despite the process clearly running (consuming CPU, writing output files at ~60 files per minute). The root cause was Python's stdout buffering behavior: when output is redirected to a file (as opposed to a terminal), Python uses block buffering (typically 8KB blocks), meaning log messages are accumulated in memory until the buffer fills or the process exits. For a multi-hour extraction job, this meant the log would remain empty for the entire duration. The fix was surgical: kill the process, add -u (unbuffered output) to the Python invocation, and restart. The extraction had already processed about 137 samples, but the script's idempotent skip logic meant no work was lost.

The extraction completed successfully, producing 17.3 million tokens of hidden states across 924 GB of data, saved in five shard directories (rows_0-2000 through rows_8000-10000). Zero errors were reported. The old vLLM-extracted hidden states (828 GB) were deleted to free space, and the original deepseek_v2.py was restored to remove the hidden state dump patch.

Training the EAGLE-3 Drafter from Scratch

With the training data in hand, the assistant turned to training a new EAGLE-3 drafter from scratch — not finetuning from the AQ-MedAI checkpoint used in earlier attempts. This was a deliberate strategic decision: the previous drafter, trained on vLLM-extracted hidden states, had achieved only ~25% acceptance rate in speculative decoding, rendering it effectively useless. The hypothesis was that SGLang-extracted hidden states, being from the actual inference engine that would be used in production, would better capture the distribution the drafter needs to predict.

The training configuration used 5 epochs over 10K samples, a learning rate of 3e-5, sequence packing at max length 2048, 3 TTT (Test-Time Training) steps, noise standard deviation 0.05, cosine scheduler with 1% warmup, and 4 data-loading worker processes. The drafter had 2.6 billion total parameters (1.19 billion trainable), using a reduced 32K draft vocabulary mapped from the verifier's 163K vocabulary — a design choice that kept the lm_head manageable at 230M parameters instead of 1.17B.

The training launched successfully, with GPU utilization at 98% and the process consuming CPU time. But then the log file went silent.

The Silent Logger: A Debugging Odyssey

The most protracted debugging challenge in this segment was the silent logger problem. The speculators library (v0.3.0), which provides the EAGLE-3 training infrastructure, uses Python's logging module with a logger named speculators.metrics. However, it never configures a logging handler. By default, Python loggers without handlers silently discard all messages below the WARNING level. This meant that the training loop was computing loss and accuracy metrics at every step — 32,304 metric events across three epochs — but none of them were being written to any log file or console.

The user's question at [msg 3453] — "what loss at first eval?" — exposed this problem directly. The assistant investigated methodically: checking process status (ps aux), confirming the process was alive with 10+ minutes of CPU time, using strace to verify that worker processes were actively reading data files, examining the speculators source code to understand the logging architecture, and confirming the training was genuinely progressing by waiting for the first epoch checkpoint to appear — a 4.7 GB model.safetensors file that materialized after 34 minutes.

The solution was surgical: kill the running process, add logging.basicConfig(level=logging.INFO) to the training script, enable checkpoint resumption via resume_from_checkpoint=True, and restart. The new log file (train_v2.log) immediately began capturing metric output. The first visible metrics showed step 0 accuracy of ~77% — a dramatic improvement over the previous drafter's 25% acceptance rate.

The Moment of Proof: Visualizing Training Progress

With metrics now flowing, the user requested charts: "create loss/acc charts so far and save here in ./train-progress/.svg/.png." The assistant extracted 32,304 lines of metric data from the remote log, wrote a custom Python plotting script (plot_metrics.py), and generated loss and accuracy charts in both PNG and SVG formats.

The numbers told a compelling story. Across three completed validation epochs:

Themes and Engineering Lessons

Several themes emerge from this segment's work that are worth articulating explicitly:

1. The primacy of observability. The silent logger problem consumed more debugging effort than any other single issue. A system that appears to be working (GPU at 98%, checkpoints saving, processes alive) but produces no visible metrics is a system you cannot trust. The fix — a single logging.basicConfig() call — was trivial once the root cause was identified, but identifying it required tracing through library source code, understanding Python's logging architecture, and systematically ruling out alternative explanations.

2. The compounding effect of infrastructure decisions. Every layer of the pipeline introduced constraints that propagated upward. The NCCL tuning affected extraction throughput. The radix cache design affected data completeness. The Python buffering behavior affected monitoring. The logging library design affected visibility. Each of these decisions, made independently, had to be understood and managed holistically.

3. The value of non-invasive modifications. The hidden state extraction patch was designed to leave no footprint on the server's normal operation — no modified return signatures, no additional dependencies, no configuration changes beyond environment variables. This design philosophy paid dividends when the extraction completed without errors and the server could be cleanly restored to its original state.

4. The necessity of verification at every step. Every significant action in this segment was followed by a verification step: benchmarking after tuning, testing extraction after patching, checking log output after launching, confirming GPU memory release after cleanup. This discipline prevented cascading failures and provided a clear audit trail when something went wrong.

5. Training from scratch can outperform finetuning from a mismatched base. The previous drafter was finetuned from AQ-MedAI's K2 drafter, which was trained on a different model (Kimi-K2 BF16 vs K2.5 INT4). Training from scratch with the correct hidden states produced dramatically better results — a lesson that challenges the assumption that transfer learning always helps.

Conclusion

This segment of the opencode session represents a complete pipeline lifecycle: from performance tuning an inference engine, through developing a custom extraction mechanism, orchestrating a multi-terabyte data extraction, training a neural network from scratch, debugging infrastructure failures, and finally validating results through visualization. The work spanned the full stack — from NCCL environment variables to Python logging architecture to matplotlib chart generation — and required deep knowledge at every layer.

The outcome was a new EAGLE-3 drafter that, by every available metric, was dramatically better than its predecessor. The 74% step 0 accuracy, the decreasing loss curves, the clean convergence across epochs — all of it validated the strategic decisions made throughout the segment: the pivot to SGLang, the non-invasive patch design, the radix cache fix, the from-scratch training approach, and the logging infrastructure repair. In the high-stakes world of large language model optimization, where a 10% improvement in speculative decoding acceptance can translate to a 2× throughput gain, this segment's work moved the project from a failing approach to a promising one.

The journey from 90 tok/s to a trained EAGLE-3 drafter was not a straight line, but it was a methodical one — and that methodical approach is what made the difference between a broken pipeline and a working one.

References

[1] "From 90 Tok/s to EAGLE-3: Engineering a Speculative Decoding Pipeline Across a Trillion-Parameter Model" — Chunk 0 article covering the full arc of Segment 25.

[2] "From SGLang Tuning to EAGLE-3 Training: A Comprehensive Pipeline Odyssey" — Chunk 1 article providing an alternative synthesis of the segment's work.