The Moment of Truth: Probing the AQ-MedAI K2 Drafter on Kimi-K2.5

In the high-stakes world of speculative decoding for large language models, few moments are as tense as the first time you load an unfamiliar drafter model against a target model it was never trained for. Message [msg 4962] captures exactly this inflection point: the server has just finished loading, the CUDA graphs have been captured, and the assistant is about to discover whether a pre-trained K2 drafter can accelerate a K2.5 model.

The Context: A Search for Better Speculation

To understand why this message matters, we need to step back. The team had been engaged in a multi-week effort to deploy the Kimi-K2.5-INT4 model (a 163,840-vocabulary, 7,168-hidden-size MoE architecture) with EAGLE-3 speculative decoding on a machine with 8 RTX PRO 6000 Blackwell GPUs. Their from-scratch EAGLE-3 drafter, trained on 100K samples of K2.5 data, had achieved a respectable 74.7% validation accuracy and ~60 tok/s throughput. But the baseline (no speculation) was running at 82 tok/s — meaning their drafter was actually slower than not using speculation at all.

The bottleneck had been traced to the verify step: approximately 30ms per cycle, of which ~25ms was consumed by 122 NCCL all-reduce operations across the 8 GPUs. This was a PCIe communication problem, not a compute problem. The team had already explored and abandoned several optimization paths — n-gram speculation was worse, NCCL_ALGO=Tree failed during CUDA graph capture, and the FlashInfer allreduce fusion fix for SM120 (Blackwell) had been applied but not yet benchmarked.

Enter the AQ-MedAI K2 EAGLE-3 drafter. This pre-trained model had been trained on 1.4 million samples of Kimi-K2 data and reportedly achieved an accept_len of 3.2–3.5 on K2. The hypothesis was straightforward: since K2 and K2.5 share the same DeepSeek V3 / MLA architecture, hidden_size (7168), and layer structure, the K2 drafter's weights might serve as either a strong initialization for fine-tuning or — in the best case — work directly out of the box.

What Happens in This Message

Message [msg 4962] is deceptively brief. The assistant writes:

Server is ready. Let me check the startup log for any warnings, then run the benchmark.

It then executes a bash command to tail the last 30 lines of the server log, revealing the tail end of CUDA graph capture — the final stage of server initialization where SGLang compiles GPU kernels into reusable execution graphs for maximum performance.

This message represents the culmination of approximately 10 minutes of waiting. The server startup sequence is complex: first, the 8-GPU model loads 64 safetensors shards of the K2.5 base model (~140GB of weights), then the AQ-MedAI drafter model (~2.4GB) is loaded, and finally CUDA graphs are captured for both the draft generation and draft extension paths. The log output shows this final stage completing with healthy metrics: "Capture draft cuda graph end. Time elapsed: 2.24 s. mem usage=0.25 GB. avail mem=8.21 GB."

The Reasoning and Motivation

The assistant's motivation in this message is methodical and risk-aware. After spending significant effort setting up the server — killing old processes, verifying GPU memory is clean, fixing the AQ-MedAI config (updating max_position_embeddings from 131072 to 262144 to match K2.5's context length), confirming the safetensors contain the required d2t vocab mapping tensor, and waiting through the long startup — the first action upon readiness is not to blindly benchmark, but to check the logs for warnings.

This is a crucial engineering discipline. A server that responds to health checks might still have issues: silent weight loading errors, graph capture failures on some GPUs, or configuration mismatches that degrade performance without crashing. By inspecting the log output first, the assistant ensures the benchmark results will be interpretable. If the server had warnings, the benchmark would need to be interpreted with caution or the issues fixed first.

The assistant also has a specific next action queued: running the benchmark script (benchmark_eagle3.py) that was copied to the server in a previous step. This script measures throughput (tok/s) across multiple runs with warmup, providing statistically meaningful results.

Input Knowledge Required

To fully understand this message, one needs to know several things:

  1. The server architecture: SGLang with 8-GPU tensor parallelism, using --tp-size 8 and --disable-custom-all-reduce (which disables the custom all-reduce kernel that is incompatible with Blackwell GPUs). The server uses --speculative-algorithm EAGLE3 with --speculative-eagle-topk 1, --speculative-num-draft-tokens 3, and --speculative-num-steps 2.
  2. CUDA graph capture: A performance optimization where GPU operations are recorded and replayed as a single fused kernel, avoiding kernel launch overhead. This takes minutes and is the last step before the server becomes ready.
  3. The AQ-MedAI drafter structure: A LlamaForCausalLMEagle3 architecture with draft_vocab_size=32000, a single transformer layer (midlayer), and vocab mapping tensors (d2t and t2d) that translate between the 32,000-token draft vocabulary and the 163,840-token target vocabulary.
  4. The vocab mapping mismatch: Earlier analysis ([msg 4948]) revealed that 31,748 out of 32,000 d2t values differ between the AQ-MedAI drafter and the team's from-scratch drafter. This is because each drafter's vocabulary mapping is built from the most frequent tokens in its training data — K2 vs K2.5 have different token distributions.

The Critical Assumption

The central assumption in this message is that the AQ-MedAI K2 drafter is architecturally compatible with K2.5. This had been verified earlier by comparing configs: both use hidden_size=7168, intermediate_size=18432, head_dim=128, and the same eagle_config with eagle_aux_hidden_state_layer_ids=[2, 30, 58]. The safetensors contain the same weight keys in the expected midlayer.* format that SGLang's llama_eagle3.py loader expects.

However, there is a deeper assumption being tested: that the hidden state representations learned by the K2 drafter are sufficiently correlated with K2.5's hidden states that the drafter can predict useful tokens. The two models share architecture but were trained on different data distributions. The K2 drafter's lm_head and midlayer weights learned to predict tokens based on K2's hidden state patterns. If K2.5 produces significantly different hidden states (due to different training data, quantization differences, or the INT4 weight representation), the drafter's predictions could be random.

What the Message Creates

This message produces the benchmark data that will determine the next several hours of work. The assistant doesn't yet know the result — that comes in the following messages ([msg 4964] onward). But the structure is set: a clean server state, a ready benchmark script, and a clear evaluation protocol.

The output knowledge created by this message includes:

The Thinking Process

The assistant's reasoning in this message reveals a careful, systematic approach. After the long wait for server startup (with the user prompting "poll health in a loop" in [msg 4960] to get faster feedback), the assistant doesn't rush. It first verifies the server state by reading the log, then proceeds to benchmark.

The log output itself is informative: it shows the final CUDA graph capture timings across multiple GPU ranks (TP0, TP1, TP2). The "avail mem=8.21 GB" line is particularly important — it confirms each GPU has sufficient free memory after loading both the base model and drafter, which is critical for stable inference.

The message also implicitly acknowledges the stakes. The team had invested significant effort in the from-scratch drafter, only to find it slower than baseline. The AQ-MedAI drafter represents a new hope: if it works well out of the box, the fine-tuning plan can be skipped entirely. If it works moderately well (accept_len > 1.0), fine-tuning is viable. If it fails completely (accept_len ~ 1.0), the entire approach needs reconsideration.

What Follows

The next messages reveal the outcome. The benchmark ([msg 4964]) shows ~52 tok/s — significantly worse than both the from-scratch drafter (60 tok/s) and baseline (82 tok/s). The accept_len is ~1.5, meaning the drafter predicts only 0.5 bonus tokens per step beyond the guaranteed first token. This is above 1.0, confirming the representations share meaningful structure, but the gap is substantial.

The assistant interprets this as a "positive signal for fine-tuning" ([msg 4966]), noting that accept_len > 1.0 means the K2 drafter's hidden states are correlated with K2.5's, just not closely enough for direct use. This launches Phase 1: fine-tuning the K2 drafter on K2.5 data — a path that will ultimately prove unsuccessful, leading to the deeper system-level optimization of the verify step that becomes the session's main theme.

Conclusion

Message [msg 4962] is a quiet but pivotal moment in a complex engineering investigation. It represents the transition from setup to measurement, from preparation to discovery. The assistant's disciplined approach — verifying server health before benchmarking — ensures that the results that follow are trustworthy. And those results, while disappointing, provide crucial information that redirects the team's efforts toward the real bottleneck: the PCIe-bound NCCL all-reduce communication that dominates the verify step. In the end, the AQ-MedAI drafter probe serves not as a solution but as a diagnostic tool, confirming that data-centric improvements alone cannot overcome the system-level communication bottleneck that defines this particular hardware configuration.