From Broken Wiring to Breakthrough: The Systematic Optimization of EAGLE-3 Speculative Decoding on Blackwell GPUs

Introduction

In the high-stakes world of large language model inference, every token per second counts. When deploying a 236-billion-parameter Mixture-of-Experts model like Kimi-K2.5 across 8 NVIDIA RTX PRO 6000 Blackwell GPUs, the difference between a deployment that barely keeps up with demand and one that comfortably exceeds baseline performance can come down to a single configuration parameter — or, more often, a chain of discoveries that transforms a broken pipeline into a working one. This article synthesizes a remarkable optimization journey spanning a single segment of an opencode coding session [1][2], in which an AI assistant and a user systematically debugged, profiled, and optimized an EAGLE-3 speculative decoding deployment from a 20% performance deficit to a 5.9% speedup over the baseline — and then identified the strategic path to even larger gains.

The journey unfolds in three distinct phases: root cause discovery and correction (fixing a hidden state wiring bug that had crippled the acceptance rate), profiling-driven optimization (adding instrumentation, tuning NCCL, and sweeping step counts to find the optimal configuration), and strategic comparison (benchmarking against a reference implementation to identify training data scale as the highest-leverage remaining improvement). Each phase builds on the previous one, and together they form a masterclass in systematic ML engineering.

Phase 1: The Hidden State Wiring Bug — When a "Fix" Breaks Everything

The story begins with a critical misunderstanding. Earlier in the session, the assistant had identified what it believed was a root cause for poor speculative decoding performance: it thought the training data used embedding-layer hidden states concatenated with layer outputs, and that the SGLang configuration needed to capture the embedding output (layer index -1). Acting on this analysis, the assistant applied a "fix," changing eagle_aux_hidden_state_layer_ids from [2, 30, 58] to [-1, 2, 30] and adding embedding capture code to the model.

This "fix" was catastrophically wrong. As the assistant later discovered by carefully tracing through the training data pipeline, the training data had never captured the embedding output. The hidden state dump patch captured at layers 3, 31, and 59 — which are the outputs of transformer layers 2, 30, and 58 respectively. The extraction script built the hidden states list as [aux_0, aux_1, aux_2, final], and the standardize_data_v1 function then concatenated cat(hs[:-1]) — meaning cat([layer3_out, layer31_out, layer59_out]). The original configuration [2, 30, 58] was correct all along.

The Moment of Discovery

The breakthrough came when the assistant wrote a norm-comparison script to compare hidden state values between training data and inference captures. By examining the first five values of each hidden state tensor, a devastating pattern emerged. The training data's hs[0] — labeled "embed" — had first-five values of [0.0295, -0.0114, -0.0170, -0.0179, -0.0183]. These values exactly matched what SGLang was capturing at "layer 3" during inference. The training data's hs[1] — labeled "layer3" — matched inference's "layer 31" capture. The layers were not off by one in the direction the assistant had assumed; they were off in the opposite direction.

Tracing through the actual code revealed the truth. The hidden state dump patch (apply_hs_dump_patch_v2.py) captured at layers 3, 31, and 59 in the layer loop — the outputs of layers 2, 30, and 58 respectively. It did not capture the embedding at all. The training extraction script (02b_extract_hidden_states_sglang.py) built the hidden states list as [aux_0, aux_1, aux_2, final] where aux_0 = layer 3 output, aux_1 = layer 31 output, aux_2 = layer 59 output. Then standardize_data_v1 used cat([layer3_out, layer31_out, layer59_out]) — exactly what the original [2, 30, 58] configuration specified.

The original config had been correct all along. The "fix" had broken it.

The Correction and Its Verification

The assistant immediately reverted the configuration, changing eagle_aux_hidden_state_layer_ids back to [2, 30, 58], killed the server, and restarted with the corrected settings. After waiting for the massive 8-GPU model to load — a process that took roughly four minutes — the assistant sent a smoke test request asking "What is 2+2? Answer briefly." The server responded correctly with "4."

But the real test was the acceptance rate. The assistant grepped the server logs for "accept_len" and "Decode batch" entries. The initial grep returned ambiguous data — the server args line happened to contain the word "accept" — but a follow-up with a longer generation request revealed the truth: the acceptance rate had jumped from ~19% to ~47%, and the acceptance length had increased from ~1.12 to ~2.85. The fix was confirmed.

This moment of verification was the culmination of a multi-hour debugging session that involved writing norm-comparison scripts, tracing through SGLang's source code, understanding the HS dump patch's capture points, reconstructing the exact data format used during training, and reverting a previous incorrect fix that had made things worse. The user's forward-thinking question — "does that impact just the model config and it will be usable in unmodded sglang, or will we need to retrain later to work with upstream sglang?" — showed concern about production deployment. The assistant confirmed that no retraining was needed: the drafter checkpoint was fully compatible with upstream SGLang, requiring only a small ~20-line patch to the model wrapper for delegation methods.## Phase 2: Profiling-Driven Optimization

With the fundamental wiring issue resolved, the assistant turned to the question of performance. The initial benchmarks were sobering: even with the correct hidden states, the EAGLE-3 speculative decoder was achieving only 71.3 tok/s — well below the ~90 tok/s baseline. Something was consuming time, but what?

The answer came from profiling instrumentation. The assistant added an EAGLE3_PROFILE=1 flag to the eagle worker that measured per-phase timing over hundreds of cycles. The results were stark: the target model verify forward pass consumed 95%+ of the cycle time (21–28ms), while the draft model was negligible at under 5%. The bottleneck was not the draft model's quality or speed, but the cost of running the full 61-layer target model forward pass to verify the draft tokens.

This discovery immediately reframed the optimization problem. Improving the draft model would yield almost no benefit. The real leverage point was reducing target verify latency. And the dominant component of that latency was allreduce communication across the 8 GPUs — every forward pass through the target model requires synchronizing gradients across all GPUs for each of the 61 transformer layers.

NCCL Tuning

The assistant then made a critical discovery: the NCCL (NVIDIA Collective Communications Library) environment variables were not configured optimally. Without NCCL tuning, the baseline server achieved only 62.9 tok/s. With NCCL tuning — specifically NCCL_PROTO=LL (Low Latency protocol), NCCL_ALGO=Ring (Ring allreduce algorithm), and NCCL_P2P_LEVEL=SYS (System-level peer-to-peer) — the baseline jumped to 88.8 tok/s, a 41% improvement.

The NCCL tuning proved equally essential for the speculative decoding server. Applying the same environment variables to the EAGLE-3 configuration reduced target verify time by approximately 27% — from 28.7ms to 21.7ms for the 5-step configuration. This was the first major optimization breakthrough: the NCCL settings reduced allreduce latency by approximately 4.6ms per token, directly attacking the dominant bottleneck.

The assistant experimented with three NCCL environment variables:

The Step Count Sweep

With NCCL tuning in place, the assistant turned to the most important architectural parameter: the number of speculative steps. The step count determines how many draft tokens are generated per cycle, which in turn determines both the verify cost and the number of accepted tokens.

The assistant's reasoning was guided by a crucial insight from the profiling data. Comparing verify times for different numbers of draft tokens revealed a surprising pattern: verifying 6 tokens took 28.7ms, while verifying 3 tokens took 25.6ms — only an 11% reduction for halving the token count. This strongly suggested that the verify cost was dominated by fixed overhead (CUDA graph replay, NCCL allreduce latency, kernel launch overhead) rather than per-token compute. If this was true, then reducing the number of steps would reduce cycle time without proportionally reducing accepted tokens.

The assistant systematically tested step counts from 1 to 10. The results formed a clear pattern:

| Config | tok/s (avg) | vs Baseline | |--------|:-----------:|:-----------:| | Baseline (NCCL tuned) | 88.8 | — | | EAGLE3 1-step | 85.1 | -4.2% | | EAGLE3 2-step | 94.0 | +5.9% | | EAGLE3 3-step | ~92 | ~+3.6% | | EAGLE3 5-step | 86.7 | -2.4% |

The 2-step configuration (producing 3 draft tokens) was the clear winner, achieving 94.0 tok/s average with peaks at 96.9 tok/s — a 5.9% improvement over the 88.8 tok/s baseline. This was the first time in the entire optimization journey that speculative decoding consistently outperformed the non-speculative baseline.

The profiling data for the optimal configuration told a precise story. Each cycle took 19.85ms total: 0.89ms for draft steps (4.5%), 18.67ms for target verify (94.1%), and 0.28ms for draft re-extend (1.4%). The accept length was approximately 1.73 tokens per cycle with an accept rate of 0.57. The 2-step configuration achieved the best ratio of accepted tokens to verify overhead — enough draft tokens to get good acceptance, but not so many that the fixed verify cost became a burden.

The assistant also performed a cost-benefit analysis that revealed the underlying economics. With baseline decode time at 11.3ms per token, and each additional draft token costing approximately 1.1ms in verify time, the expected net benefit of each step was positive: approximately 7.4ms saved for the first step, declining to 6.0ms for the fifth. Yet the empirical data showed that 5 steps underperformed 2 steps. The discrepancy pointed to complexities — CUDA graph overhead, memory bandwidth contention, scheduler behavior — that only emerged in actual measurement.

The assistant also noted a critical interaction between num_steps and num_draft_tokens. In SGLang's EAGLE-3 implementation, when topk=1, the server enforces num_draft_tokens = num_steps + 1. This means that specifying --speculative-num-steps 2 produces 3 draft tokens — the initial token plus two speculative steps. Understanding this constraint was essential for interpreting the sweep results correctly.

Phase 3: The Strategic Pivot — Comparative Analysis

Just as the assistant was celebrating the 94 tok/s breakthrough, the user interjected with a question that would reframe the entire project: "How does our model size and dataset compare to AQ-MedAI's Kimi-K2-Instruct-eagle3? Seems that one is 1B and ours is 2B-ish?"

This question arrived at a critical inflection point. The engineering optimizations had pushed the system to its limits — NCCL tuning and step count optimization had extracted a 5.9% improvement, but the accept length of ~2.1 was still far below the theoretical potential of speculative decoding. The user's question implicitly asked: "Are we optimizing the right thing?"

The assistant's investigation produced a detailed side-by-side comparison that revealed a stunning finding: the architectures were identical. Both models used hidden_size=7168, intermediate_size=18432, num_attention_heads=64, num_key_value_heads=64, and critically, the same eagle_layer_ids=[2, 30, 58]. The apparent parameter count discrepancy (1B vs 2.6B) was entirely an artifact of counting conventions — the AQ-MedAI team likely excluded the frozen embed_tokens (1.17B parameters) and lm_head (229M parameters) from their count, since those weights are shared with and frozen from the target model. The trainable parameters were nearly identical.

With architecture ruled out, the assistant compared the remaining variables. The numbers were staggering:

The Results: What Was Achieved

The optimization journey produced several concrete outcomes that can be quantified and celebrated:

A verified optimal configuration: 2 speculative steps (3 draft tokens) with NCCL tuning (NCCL_PROTO=LL, NCCL_ALGO=Ring, NCCL_P2P_LEVEL=SYS) achieving 94 tok/s.

A 5.9% speedup over baseline: From 88.8 tok/s to 94.0 tok/s, with peaks at 96.9 tok/s. This was the first time in the entire project that speculative decoding consistently outperformed the non-speculative baseline.

A precise performance breakdown: Draft steps 0.89ms (4.5%), Target verify 18.67ms (94.1%), Draft re-extend 0.28ms (1.4%), total cycle 19.85ms. This breakdown provides a clear target for future optimization: any improvement to the target verify time directly translates to throughput gains.

A quantified NCCL tuning impact: 27% reduction in verify time, from 28.7ms to 18.67ms. This was achieved purely through environment variable configuration, without any code changes.

A strategic direction: Training data expansion identified as the highest-leverage remaining improvement, with a 38× data gap to close. The comparison against AQ-MedAI's model showed that accept lengths of 3.2-3.5 are achievable with sufficient training data, which would translate to even larger speedups.

A cleaned-up codebase: The debug logging and instrumentation added during the investigation were cleaned up in a deliberate transition from debugging mode to optimization mode. The debug prints prevented CUDA graph capture and distorted performance measurements, so removing them was a prerequisite for accurate benchmarking.

The Broader Implications

This segment's journey has implications beyond the specific deployment of Kimi-K2.5 on Blackwell GPUs. Several lessons generalize to other speculative decoding deployments:

Speculative decoding is not a free lunch. The common assumption that adding a draft model will always improve throughput is false. The draft model adds overhead — the target model must verify each draft token — and that overhead can outweigh the benefits if the draft model is inaccurate or the verification cost is high. In this case, the verification cost was dominated by allreduce communication across 8 GPUs, which is a common pattern in multi-GPU deployments.

The optimal step count depends on the specific hardware and model configuration. There is no universal "best" step count. The assistant's sweep from 1 to 10 steps revealed that 2 steps was optimal for this configuration, but a different hardware setup (e.g., with NVLink instead of PCIe) or a different model (e.g., with fewer layers or faster attention) might have a different optimum. The only way to find the optimum is to measure empirically.

NCCL tuning is a high-leverage optimization for multi-GPU inference. The 27% reduction in verify time from NCCL tuning was the single largest optimization gain in this segment. Many practitioners overlook NCCL environment variables, assuming that the defaults are optimal. This segment demonstrates that NCCL tuning can have a dramatic impact, especially in PCIe-only configurations where allreduce is the dominant bottleneck.

Training data quality and quantity matter more than architectural optimization. The comparison against AQ-MedAI's model revealed that the architectures were identical, but the training data volumes differed by 38×. The performance gap (accept length 3.2-3.5 vs ~2.1) was directly attributable to this data gap. This is a reminder that in ML systems, the data often matters more than the model architecture or inference optimization.

The Path Forward

The assistant's work in this segment has laid a solid foundation for future improvements. The immediate next steps are clear:

Scale up training data generation. The 38× data gap identified in the comparative analysis is the single highest-leverage improvement. The team has already built a robust pipeline for generating synthetic training data using the target model itself — the hidden state extraction, standardization, and training loop are all functional. The only missing piece is more data. Generating 1.4 million samples (or even a fraction of that) would likely produce a drafter with accept lengths of 3.0+, which would translate to throughput gains of 20-30% or more.

Retrain the drafter with expanded data. Once sufficient training data is available, the drafter should be retrained from scratch. The training pipeline is already validated — it was used to train the current 37K-sample model. The same pipeline can be used with a larger dataset, and the training time should scale linearly with the dataset size.

Re-evaluate the optimal step count with the improved drafter. A more accurate drafter may shift the optimal step count. With higher acceptance rates, deeper speculation (3-4 steps) might become viable, potentially yielding even larger speedups. The assistant should re-sweep the step count parameter after retraining.

Consider upstream compatibility. The user's question about compatibility with unmodified SGLang is important. The current deployment requires a small ~20-line patch to the model wrapper for delegation methods. If the team wants to deploy with upstream SGLang, they should either contribute this patch upstream or find a way to make the drafter compatible without patches.

Conclusion

This segment of the opencode session tells a story that every ML engineer will recognize: the journey from a broken pipeline to a working one, from guesswork to measurement, from tactical optimization to strategic insight. The assistant and user together debugged a subtle hidden state wiring error, added profiling instrumentation to understand where time was actually going, tuned the communication layer to reduce allreduce overhead, swept the configuration space to find the optimal step count, and then stepped back to compare against an external reference — discovering that the real leverage lay not in further engineering optimization but in scaling up training data.

The 5.9% speedup over baseline is real and meaningful, but it is not the end of the story. It is the foundation upon which larger gains will be built. The assistant has transformed speculative decoding from a theoretical concept into a measured reality, and in doing so has created a clear roadmap for the next phase of work: generate more training data, retrain the drafter, and close the gap with the state of the art. The systematic methodology demonstrated here — instrument, measure, isolate, sweep, compare — is a template that can be applied to any ML inference optimization challenge.

The key insights from this segment are worth restating: