The Architecture Check: Confirming Drop-In Compatibility Between Kimi-K2 and K2.5 Draft Models

In the sprawling effort to accelerate inference of the 1-trillion-parameter Kimi-K2.5 model through speculative decoding, a single message at index 4925 represents a quiet but pivotal turning point. After days of debugging, profiling, and tuning — after discovering that EAGLE-3 speculation was actually hurting performance (59–61 tok/s vs. a baseline of 82–83 tok/s), after confirming that the verify step's 30ms cycle time was an irreducible cost of running without CUDA graphs on 8 PCIe GPUs — the assistant arrives at a strategic fork in the road. The message reads, in its entirety:

Now let me compare the two configs side by side: [bash] ssh root@10.1.230.174 'cat /data/eagle3/output_100k_sglang/4/config.json' { "architectures": [ "LlamaForCausalLMEagle3" ], "attention_bias": false, "attention_dropout": 0.0, "bos_token_id": 163584, "draft_vocab_size": 32000, "dtype": "bfloat16", "eagle_config": { "eagle_aux_hidden_state_layer_ids": [ 2, 30, 58 ], "use_aux_hidden_state": true }, "eos_token_id": 163585, "head_dim": 128, "hidden_act": "silu", "hidden_size": 7168, "initializer_range": 0.02, "intermediate_size": 18432, "max_position_embeddin...

On its surface, this is a simple bash command reading a JSON configuration file. But in the broader narrative of the session, it is the moment the assistant pivots from chasing inference-engine optimizations to embracing a data-and-model-quality strategy — and the first step is verifying that a pre-trained draft model from a different but related model family can serve as a drop-in starting point for fine-tuning.

The Strategic Context: Why This Comparison Matters

To understand the weight of this message, one must trace the thread of frustration that preceded it. The assistant had spent multiple rounds ([msg 4900] through [msg 4919]) systematically investigating why EAGLE-3 speculative decoding was delivering only ~60 tok/s when the baseline without speculation was 82–83 tok/s. The root cause was identified: the verify step — where the target model checks the draft tokens — runs in "extend" mode, which cannot use CUDA graphs. This costs approximately 30 milliseconds per cycle, regardless of whether --speculative-attention-mode is set to prefill or decode. The assistant had tried patching engine.py, patching scheduler.py, and even persisting NCCL tuning environment variables in /usr/lib/python3.12/sitecustomize.py — none of these interventions reduced the verify time.

The math was brutal. With a 30ms cycle time and an average acceptance length of 2.0 tokens per cycle, the effective throughput was 2.0 / 0.030 = 67 tok/s before streaming overhead, which matched the observed ~60 tok/s. To merely break even with the 82 tok/s baseline, the acceptance length needed to be 2.46 tokens per cycle. To reach 150 tok/s, it needed 4.5 tokens per cycle, requiring approximately 78% conditional accuracy per draft step — far beyond the ~64% their current drafter achieved after training on 37,000 samples.

The user then asked a crucial question ([msg 4922]): "Check K2 AQ-MedAI model shape to see if we can finetune it for K2.5." This question reframes the entire problem. Instead of trying to squeeze more performance out of the inference engine — which appeared to have hit a hard wall — the user proposes improving the draft model itself by leveraging transfer learning from a model trained on 1.4 million samples.

The Decision Process Visible in the Message

The assistant's response in message 4925 is deceptively simple, but the reasoning behind it is multi-layered. The assistant had already fetched the AQ-MedAI Kimi-K2 drafter's config.json from HuggingFace in the preceding message ([msg 4924]). Now it reads its own trained drafter's configuration to perform a side-by-side structural comparison.

The key architectural parameters that must match for a successful transfer are:

1. hidden_size: 7168 — the dimensionality of the hidden states
2. intermediate_size: 18432 — the feed-forward network width
3. head_dim: 128 — the per-head dimension for attention
4. draft_vocab_size: 32000 — the vocabulary size for draft token prediction
5. num_attention_heads and num_key_value_heads — the attention head counts
6. num_hidden_layers — the depth of the transformer If these parameters match exactly between the K2 and K2.5 drafters, then the AQ-MedAI weights can be loaded directly into the K2.5 architecture, providing a vastly better initialization than the random or small-data initialization used for the current 37K-sample drafter. The AQ-MedAI drafter was trained on 1.4 million samples and achieves an acceptance length of 3.2–3.5 tokens per cycle — a 60–75% improvement over the current 2.0.

Assumptions Embedded in the Approach

The assistant makes several assumptions in this message. The first is that architectural compatibility between K2 and K2.5 drafters is sufficient for transfer learning to work. This is not guaranteed: even if the shapes match, the hidden state distributions could differ enough that fine-tuning requires more data or different hyperparameters. The K2.5 model is a fine-tuned version of K2, and while the tokenizer and base architecture are the same family, the hidden states at the auxiliary layers (layers 2, 30, and 58, as specified in the eagle_config) could have shifted enough to degrade the drafter's accuracy.

A second assumption is that the AQ-MedAI drafter's weights are publicly available and accessible. The HuggingFace search in [msg 4923] confirmed the repository exists, but the assistant has not yet verified that the actual weight files (safetensors) are downloadable or that the license permits fine-tuning.

A third assumption — perhaps the most critical — is that fine-tuning a better drafter is the highest-leverage path forward. The assistant had previously identified that reducing the verify cycle time would have an even更大 impact: if verify could be cut from 30ms to 15ms, break-even would require only 2.25 accept_len (56% accuracy, easily achievable). But the assistant has implicitly concluded that the verify time is a hard constraint of the hardware (8 PCIe GPUs with a 1T MoE model) and the SGLang framework's lack of CUDA graph support for extend-mode verification.

Input Knowledge Required

To fully understand this message, one needs substantial background knowledge spanning several domains. First, familiarity with the EAGLE-3 speculative decoding algorithm — specifically how it uses a lightweight draft model to propose multiple tokens, which the target model then verifies in a single forward pass. Second, understanding of the SGLang inference framework and its distinction between "decode" mode (which supports CUDA graphs for kernel fusion) and "extend" mode (which does not). Third, knowledge of transformer architecture parameters: what hidden_size, intermediate_size, head_dim, and draft_vocab_size mean and why they must match for weight compatibility. Fourth, awareness of the Kimi model family tree — that Kimi-K2.5 is a fine-tuned derivative of Kimi-K2, sharing the same base architecture and tokenizer. Fifth, understanding of the training data scaling relationship: that AQ-MedAI's 1.4 million samples produced 3.2–3.5 accept_len while 37,000 samples produced only 2.0.

Output Knowledge Created

This message produces several concrete pieces of knowledge. First, it confirms that the locally trained drafter uses the LlamaForCausalLMEagle3 architecture — the standard EAGLE-3 architecture for LLaMA-family models. Second, it reveals the specific auxiliary hidden state layers used (2, 30, 58), which is critical for understanding how the draft model interfaces with the target model. Third, it surfaces the exact parameter values that will be compared against the AQ-MedAI config: hidden_size=7168, intermediate_size=18432, head_dim=128, draft_vocab_size=32000. Fourth, it establishes a baseline configuration that any fine-tuned model must preserve for compatibility with the existing SGLang deployment pipeline.

The Thinking Process Visible in the Reasoning

The assistant's reasoning in this message is compressed but discernible. The phrase "Now let me compare the two configs side by side" reveals that the assistant has already retrieved the AQ-MedAI config (in the previous message) and is now retrieving the local config to perform the comparison. The assistant does not simply trust that the architectures match — it verifies empirically by reading the actual files from disk.

The choice to read the config from /data/eagle3/output_100k_sglang/4/config.json rather than from a training script or documentation is significant. It demonstrates a commitment to ground truth: the configuration that actually produced the deployed model, not the configuration that was intended or assumed. This is particularly important in a debugging context where assumptions have repeatedly been proven wrong — the 94 tok/s that wasn't reproducible, the decode attention mode that didn't help, the NCCL tuning that didn't propagate.

The message also reveals a shift in the assistant's mental model. Earlier messages were focused on the inference engine — trying different attention modes, patching environment variable propagation, checking CUDA graph status. This message shifts focus to the model itself. The assistant is no longer asking "how can we make the engine faster?" but rather "how can we make the drafter better?" This is a fundamental reframing of the problem, and this message is the first concrete step in that new direction.

The Broader Significance

In the arc of the entire coding session, message 4925 is the hinge point. Everything before it is diagnosis and dead ends; everything after it is the fine-tuning campaign. The chunk summary for segment 33 confirms that the session concluded with writing a comprehensive eagle-k2finetune-game-plan.md document covering three approaches: fine-tuning AQ-MedAI's drafter with existing K2.5 samples, scaling training data to 200K+ samples, and a direct plug-in probe to measure hidden state similarity. But none of that would be possible without first confirming architectural compatibility — which is precisely what this message accomplishes.

The message also embodies a broader lesson about debugging complex systems: when you hit a hard performance wall in the inference path, the most productive response may not be to keep hammering at the wall, but to change the problem entirely. The 30ms verify time was not going to budge without fundamental changes to SGLang's CUDA graph support or a hardware upgrade. But the draft model's accuracy was a lever that could be pulled with existing tools and data — and the AQ-MedAI drafter offered a shortcut to a much better starting point.

Conclusion

Message 4925 is, on its surface, a single bash command reading a JSON file. But in context, it is the moment a multi-day debugging effort pivots from chasing inference-engine optimizations to embracing a data-and-model-quality strategy. It is the moment the assistant stops asking "how do we make the engine faster?" and starts asking "how do we make the drafter better?" The architectural comparison it initiates will determine whether transfer learning from AQ-MedAI's Kimi-K2 drafter is feasible — and that feasibility assessment is the foundation for the entire fine-tuning campaign that follows. In the high-stakes world of 1T-parameter model inference, sometimes the most important work is not tuning kernels but reading config files.