The Pivot to AQ-MedAI: Testing an Alternative EAGLE-3 Drafter After the Hidden State Fix

Introduction

In the long arc of debugging and optimizing speculative decoding for the Kimi-K2.5 model, message <msg id=3626> represents a pivotal moment of empirical comparison. After spending hours tracing a devastating bug—the server was launched with --speculative-algorithm EAGLE instead of EAGLE3, causing the entire auxiliary hidden state capture mechanism to be silently bypassed—the assistant had finally confirmed that hidden states were arriving at the draft model as the correct 21504-dimensional concatenated vectors. The custom drafter trained on 10,000 SGLang-extracted samples was now functional, with an acceptance length of approximately 2.1 tokens per verification step. Yet the throughput remained stubbornly below the non-speculative baseline: 56.7 tok/s versus 90 tok/s.

This message captures the moment the assistant pivots to test a different drafter entirely: the AQ-MedAI K2 drafter, a pre-trained model developed by the open-source community. The reasoning is clear and data-driven: if the custom drafter's low acceptance rate stems from insufficient training data (only 10K samples), perhaps a drafter trained on a much larger corpus—even if not specifically fine-tuned on Kimi-K2.5—might achieve higher acceptance and finally deliver the speedup that speculative decoding promises.

The Message in Full

The assistant executes a single bash command via SSH to the remote server at 10.1.230.174:

ssh root@10.1.230.174 'sleep 3 && nohup bash -c "NCCL_PROTO=LL NCCL_ALGO=Ring NCCL_P2P_LEVEL=SYS NCCL_MAX_NCHANNELS=16 NCCL_BUFFSIZE=16777216 NCCL_NTHREADS=512 SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 ~/ml-env/bin/python3 -m sglang.launch_server \
  --model-path /shared/kimi-k2.5-int4 --trust-remote-code --tp-size 8 \
  --mem-fraction-static 0.85 --host 0.0.0.0 --port 8000 \
  --speculative-algorithm EAGLE3 \
  --speculative-draft-model-path /data/eagle3/aq-medai-k2-drafter \
  --speculative-num-steps 3 --speculative-eagle-topk 4 --speculative-num-draft-tokens 16 \
  --num-continuous-decode-steps 4 --disable-custom-all-reduce \
  --disable-cuda-graph \
  --log-level info" > /data/eagle3/sglang_eagle3_aqmedai.log 2>&1 &'
echo "Starting with AQ-MedAI drafter..."

The bash tool reports that the command exceeded its 15,000 ms timeout and was terminated. However, because the server launch was wrapped in nohup and backgrounded with &, the process should continue running on the remote machine regardless of the SSH client's timeout.

The Reasoning and Motivation

To understand why this message was written, we must trace the chain of reasoning that led to this moment. The assistant had just completed a series of benchmarks on the custom-trained drafter (checkpoint from output_10k_sglang/4). With --speculative-num-draft-tokens 16, the server logged an acceptance length of ~2.2 and an acceptance rate of ~0.14 (14%). With --speculative-num-draft-tokens 5, the acceptance rate improved to ~0.41 but the acceptance length remained at ~2.0-2.2. In both cases, throughput was 53-57 tok/s—far below the 90 tok/s non-speculative baseline.

The assistant's analysis in the preceding messages was precise: "The fundamental problem: accept_len of ~2.1 is too low. With an ideal EAGLE-3 drafter, accept_len should be 3-4+. Our drafter trained on only 10K samples just isn't good enough." This diagnosis frames the problem as a data quantity issue. The EAGLE-3 paper demonstrates clear scaling laws: more training data produces drafters with higher acceptance rates. Ten thousand samples, while sufficient to demonstrate the pipeline works, is far from the hundreds of thousands or millions used in the original paper.

But the assistant also knows there's another drafter available on the system: the AQ-MedAI K2 drafter. Earlier in the conversation ([msg 3624]), the assistant inspected its configuration and confirmed it has the same architecture (LlamaForCausalLMEagle3), the same hidden size (7168), the same draft vocabulary size (32000), and critically, the same auxiliary hidden state layer IDs ([2, 30, 58]). This means it is structurally compatible with the Kimi-K2.5 target model, which has 60 layers total and captures layers 2, 30, and 58 (the 1st, 30th, and 58th layers, zero-indexed).

The AQ-MedAI drafter was trained as a general K2 drafter—not specifically fine-tuned on Kimi-K2.5 outputs. This is both a strength and a weakness. It was likely trained on far more data (AQ-MedAI is a well-funded AI lab), so its fc projection layer and midlayer transformer block may be better optimized. However, it was trained on a different base model's distribution, so its predictions may not align well with Kimi-K2.5's token distribution. The assistant is about to find out which effect dominates.

Decisions Made in This Message

This message embodies several key decisions:

Decision 1: Test the AQ-MedAI drafter before scaling up training data. The assistant could have immediately launched a massive data collection pipeline to train a better custom drafter. Instead, it chose to first benchmark the existing AQ-MedAI drafter. This is a smart prioritization: if the AQ-MedAI drafter achieves 3-4× acceptance length, it would save days of data collection and retraining. If it fails, the assistant hasn't lost much—the server restart takes only a few minutes.

Decision 2: Use the same hyperparameters as the custom drafter test. The assistant launches with --speculative-num-steps 3, --speculative-eagle-topk 4, and --speculative-num-draft-tokens 16. This is the same configuration used in the first custom drafter benchmark ([msg 3607]). Keeping hyperparameters constant ensures a fair comparison—any difference in throughput can be attributed to the drafter itself, not to configuration changes.

Decision 3: Keep CUDA graphs disabled. The --disable-cuda-graph flag is carried over from previous launches. CUDA graphs would significantly improve performance by capturing GPU operations as reusable graphs, but they were disabled earlier for debugging purposes. The assistant is prioritizing correctness and comparability over peak performance.

Decision 4: Log to a separate file. The output is redirected to /data/eagle3/sglang_eagle3_aqmedai.log, keeping it separate from the custom drafter logs (sglang_eagle3_v3.log and sglang_eagle3_v3b.log). This organization makes it easy to compare results later.

Assumptions Made

The assistant makes several assumptions in this message, some explicit and some implicit:

Assumption 1: The AQ-MedAI drafter is compatible with Kimi-K2.5. The assistant verified the architecture and hidden size match, but there's a deeper assumption: that the drafter's embed_tokens and lm_head weight matrices (both of size 32000 × 7168) are compatible with Kimi-K2.5's vocabulary. The Kimi-K2.5 model uses a vocabulary of 163,840 tokens (as seen in the checkpoint dump showing embed_tokens.weight: torch.Size([163840, 7168])), but the drafter's vocabulary is only 32,000. This mismatch is actually expected—EAGLE-3 drafters typically use a reduced vocabulary of the most frequent tokens, mapping them via the d2t (draft-to-target) embedding. The assistant previously verified that the custom drafter's d2t tensor has shape [32000], which maps draft vocabulary indices to target vocabulary indices. The AQ-MedAI drafter's d2t was not inspected, but the assistant assumes it exists and is correctly shaped.

Assumption 2: The server will start successfully despite the SSH timeout. The assistant wraps the command in nohup and backgrounds it, correctly assuming the process will outlive the SSH session. The timeout of 15 seconds is a tool-level constraint, not a sign of failure.

Assumption 3: The AQ-MedAI drafter's weights will be loaded correctly by SGLang. The assistant previously discovered that the custom drafter's fc.weight had a mean of 0.000000 in bfloat16 (due to rounding), which caused a moment of panic before confirming the weights were non-zero. The assistant assumes the AQ-MedAI drafter's weights are similarly well-formed and will load without issues.

Assumption 4: The benchmark results will be interpretable. The assistant plans to run the same benchmark script (/tmp/eagle3_bench.py) after the server starts, comparing the AQ-MedAI drafter's throughput and acceptance metrics directly against the custom drafter's numbers.

Mistakes and Incorrect Assumptions

Potential mistake: Not checking the AQ-MedAI drafter's weight statistics before launching. In the previous round ([msg 3624]), the assistant inspected the AQ-MedAI drafter's config.json but did not dump its actual weight tensors. The custom drafter had a critical moment where fc.weight appeared to be all zeros (mean=0.000000), which turned out to be a bfloat16 rounding artifact. The AQ-MedAI drafter could have genuine issues—missing keys, incorrect shapes, or zero-initialized weights—that would only be discovered after the server starts and fails.

Potential mistake: Not verifying the d2t mapping. The assistant knows the custom drafter has a d2t tensor of shape [32000] that maps draft vocabulary indices to target vocabulary indices. For the AQ-MedAI drafter, this mapping might be incompatible with Kimi-K2.5's vocabulary ordering. If the mapping is wrong, the drafter would predict plausible-looking tokens that map to completely unrelated target tokens, resulting in zero acceptance. The assistant didn't check this before launching.

Assumption about training distribution: The assistant assumes that "more training data = better acceptance" and that the AQ-MedAI drafter, being trained on more data, might outperform the custom drafter. However, the AQ-MedAI drafter was trained on a different base model's distribution. If Kimi-K2.5 has a significantly different token distribution (which is likely, given it uses a 163,840-token vocabulary), the AQ-MedAI drafter's predictions might be systematically misaligned. The assistant acknowledges this implicitly by calling it "a general K2 drafter, not K2.5-specific."

Input Knowledge Required

To fully understand this message, one needs:

  1. The EAGLE-3 architecture: Understanding that EAGLE-3 uses a draft model that takes concatenated hidden states from multiple layers of the target model (here layers 2, 30, 58, producing 3 × 7168 = 21504 dimensions), passes them through an fc projection layer (21504 → 7168), then through a single transformer block (midlayer), and finally through a shared lm_head to produce draft token predictions.
  2. The speculative decoding pipeline: Understanding that --speculative-num-steps 3 means the draft model runs 3 autoregressive steps to generate candidates, --speculative-eagle-topk 4 means it samples from the top-4 logits at each step (producing a tree of candidates), and --speculative-num-draft-tokens 16 is the total number of candidate tokens in the tree. The target model then verifies all candidates in parallel using a tree attention mask.
  3. The server infrastructure: The machine has 8 RTX PRO 6000 Blackwell GPUs (SM120 architecture), uses tensor parallelism (--tp-size 8), and runs SGLang with NCCL optimizations (NCCL_PROTO=LL, NCCL_ALGO=Ring, etc.).
  4. The debugging history: The critical bug where --speculative-algorithm EAGLE was used instead of EAGLE3, causing the is_eagle3() check to fail and the auxiliary hidden state capture to never activate. This meant the draft model received 7168-dim final-layer states instead of 21504-dim concatenated states, rendering the fc layer useless.
  5. The custom drafter's performance: 56.7 tok/s average with 16 draft tokens, 53.2 tok/s with 5 draft tokens, versus 90 tok/s non-speculative baseline. Acceptance length ~2.1, acceptance rate ~14-41% depending on configuration.

Output Knowledge Created

This message creates several outputs:

  1. A new SGLang server process running with the AQ-MedAI drafter, logging to /data/eagle3/sglang_eagle3_aqmedai.log. This server will be benchmarked in subsequent messages.
  2. A direct comparison point: Once the AQ-MedAI drafter is benchmarked, the assistant will be able to answer: "Is our custom 10K-sample drafter better than a general-purpose drafter trained on more data?" This is a crucial data point for deciding whether to invest in scaling up training data or to abandon the EAGLE-3 approach entirely.
  3. Confirmation of the pipeline's modularity: By successfully swapping the draft model path from /data/eagle3/output_10k_sglang/4 to /data/eagle3/aq-medai-k2-drafter, the assistant demonstrates that the SGLang EAGLE-3 infrastructure correctly handles different draft model checkpoints, validating the architecture's plug-and-play design.

The Thinking Process Visible in Reasoning

The assistant's thinking process is visible across the sequence of messages leading to this one. The key reasoning chain is:

  1. Diagnose: The custom drafter achieves accept_len ~2.1 but throughput is below baseline. Why? Because the overhead of running the draft model + verifying 16 candidates exceeds the savings from accepting ~2 tokens per step.
  2. Hypothesize: The EAGLE-3 paper shows acceptance scales with training data. 10K samples is too little. More data would produce a better drafter.
  3. Explore alternatives: Before committing to a multi-day data collection and retraining pipeline, check if an existing drafter (AQ-MedAI) performs better. It was trained on more data, even if not K2.5-specific.
  4. Execute: Launch the AQ-MedAI drafter with identical parameters for a fair comparison. This is classic scientific method: observe a problem, form a hypothesis, design an experiment, execute it, and evaluate the results. The assistant is systematically narrowing down the cause of the performance gap, ruling out the drafter quality before investing in more data.

Conclusion

Message <msg id=3626> is a deceptively simple command that represents a critical juncture in the EAGLE-3 optimization effort. After fixing the hidden state concatenation bug and confirming that the custom drafter works but underperforms, the assistant pivots to test an alternative hypothesis: that the drafter itself, not just the training data quantity, is the limiting factor. The AQ-MedAI drafter launch is a low-cost, high-information experiment that will either validate the "more data" path or redirect efforts toward a different approach entirely.

The message embodies good engineering practice: test the simplest alternative first, control for variables, and let data drive decisions. Whether the AQ-MedAI drafter succeeds or fails, the knowledge gained from this experiment will inform the next steps in the months-long effort to make speculative decoding work for Kimi-K2.5.