The Vocab Mapping Epiphany: A Critical Decision in the EAGLE-3 Drafter Probe

In the high-stakes world of large language model inference optimization, few moments are as quietly decisive as the one captured in message [msg 4949] of this coding session. The assistant, deep in the process of probing whether a pre-trained EAGLE-3 drafter from Kimi-K2 could be directly plugged into a Kimi-K2.5 inference pipeline, made a discovery that could have derailed the entire effort — or saved it. The message is a masterclass in applied reasoning under uncertainty, where a seemingly alarming result (31,748 out of 32,000 values differing between two vocabulary mapping tensors) is correctly interpreted not as a bug, but as an expected and even necessary condition for the experiment to work at all.

The Scene: Phase 0 of the K2→K2.5 Probe

To understand the significance of this message, one must first understand the broader context. The team had been engaged in a multi-week effort to deploy the massive Kimi-K2.5 INT4 model (a 1-trillion-parameter Mixture-of-Experts architecture) on a remote machine with 8× NVIDIA RTX PRO 6000 Blackwell GPUs connected via PCIe (no NVLink). The baseline inference throughput was a respectable 82 tok/s, but the team had invested heavily in EAGLE-3 speculative decoding — a technique where a smaller "draft" model proposes tokens that a larger "target" model verifies in parallel — only to find that their from-scratch EAGLE-3 drafter was actually hurting performance, achieving only 60 tok/s.

The bottleneck had been traced to the verify step, which consumed ~30ms per cycle (97% of the total time), dominated by 122 NCCL all-reduce operations across the 8 GPUs. The team's own drafter, trained on only 37K samples, achieved an accept length of ~2.0 tokens — well below the break-even point of 2.46 needed to overcome the verify overhead.

Enter the AQ-MedAI K2 EAGLE-3 drafter. This pre-trained model, trained on 1.4 million samples and achieving accept lengths of 3.2–3.5 on Kimi-K2, seemed like a perfect candidate for fine-tuning. Both models share the same DeepSeek V3 / MLA architecture, same hidden_size (7168), same layer structure — the weights should be an ideal initialization. The game plan called for three phases: Phase 0 (a quick probe, just plugging in the K2 drafter with K2.5 to measure hidden state similarity, ~30 minutes), Phase 1 (fine-tuning on existing 37K samples, ~2-3 hours), and Phase 2 (scaling up training data, ~2-3 days).

Message [msg 4949] occurs during Phase 0. The assistant has already killed existing server processes, verified the AQ-MedAI drafter directory, confirmed the GPUs are clean, and inspected the safetensors file. Now it's comparing the vocabulary mappings between the AQ-MedAI drafter and the team's own drafter.

The Discovery: 31,748 Differences

The assistant runs a Python script via SSH that loads both d2t (draft-to-target) tensors from the two safetensors files and compares them. The result is stark:

d2t match: False
  31748 / 32000 values differ

Only 252 out of 32,000 token positions are the same. On the surface, this looks catastrophic. The d2t mapping is how the drafter's output logits (in a 32,000-token "draft vocab") get mapped back to the target model's actual vocabulary (which for Kimi-K2.5 is 163,840 tokens). If the mappings are completely different, how could the drafter possibly work?

The assistant's initial reaction is measured: "The vocab mappings are different — 31,748 out of 32,000 values differ. This is expected since AQ-MedAI built their mapping from K2 training data while ours was built from K2.5 data."

But then comes the crucial moment of deeper reasoning. The assistant catches itself and writes:

"Wait — actually this is a significant issue."

This self-correction is the heart of the message. The assistant realizes that the lm_head layer of the drafter — the final classification layer that produces logits over the draft vocabulary — was trained with AQ-MedAI's specific d2t mapping. The neural network learned to predict specific positions in the draft vocab that correspond, via that specific mapping, to particular target tokens. Using a different mapping would be like giving someone a map where all the street names have been changed — the directions (logits) point to the right places on the old map but lead to completely wrong destinations on the new one.

The Correct Decision: Use AQ-MedAI's Mapping

The assistant correctly concludes: "for the direct probe, we MUST use AQ-MedAI's d2t (already in their safetensors)."

This is the right call, and it's grounded in a deep understanding of how SGLang's EAGLE-3 implementation works. The assistant had previously verified (in [msg 4948]) that SGLang loads d2t directly from the draft model's safetensors and creates hot_token_id = d2t + arange(32000) — converting from a delta format to absolute token IDs. The t2d (target-to-draft) mapping is skipped entirely. This means the AQ-MedAI safetensors file is self-contained: it carries its own vocabulary mapping that the lm_head was trained against.

The decision has a cascading implication: if Phase 0 shows promise and the team proceeds to Phase 1 (fine-tuning), they will need to decide whether to keep AQ-MedAI's mapping or switch to the K2.5 mapping. Keeping the mapping means the lm_head continues to predict tokens in the same "coordinate system" it was trained on, which should accelerate fine-tuning. Switching would require the model to relearn the mapping from scratch, potentially negating the benefit of the pre-trained initialization.

The Config Fix: A Practical Detail

Having resolved the conceptual question, the assistant executes a practical fix: updating the max_position_embeddings in the AQ-MedAI config from 131072 to 262144. This is necessary because Kimi-K2.5 supports a longer context window than K2, and SGLang will reject the config if the value is too low. The assistant reads the config, modifies it in Python, and writes it back — a straightforward operation but one that would cause a cryptic error if forgotten.

The assistant also takes care to preserve all other config fields, including the critical eagle_config with eagle_aux_hidden_state_layer_ids: [2, 30, 58] — the three layers from which hidden states are extracted for the drafter. These layer IDs had been a source of bugs earlier in the session (the config was initially [-1, 2, 30], which caused a hidden state mismatch that took days to diagnose).

Assumptions and Knowledge Required

To fully understand this message, several pieces of background knowledge are essential:

  1. EAGLE-3 architecture: The drafter uses a "draft vocabulary" of 32,000 tokens, which is a subset or remapping of the target model's full vocabulary (163,840 tokens for Kimi-K2.5). The d2t tensor maps each of the 32,000 draft token positions to a target token ID. The lm_head produces logits over the draft vocabulary, and the mapping converts these to actual token predictions.
  2. SGLang's loading mechanism: The inference engine loads the d2t tensor from the drafter's safetensors file and uses it to construct hot_token_id — the set of target token IDs that the drafter is allowed to predict. This happens automatically; there is no separate configuration file for the mapping.
  3. The difference between K2 and K2.5 vocabularies: Although both models share the same DeepSeek V3 / MLA architecture, they have different vocabularies (different tokenizer training data). The top-32,000 most frequent tokens in K2 training data will differ substantially from the top-32,000 in K2.5 data, which explains why only 252 positions match.
  4. The concept of "accept length": In speculative decoding, the draft model proposes N tokens, and the target model verifies them in parallel. The "accept length" is the average number of tokens accepted per verification cycle. Higher accept lengths mean more tokens generated per target model invocation, which is critical for overcoming the verify overhead.
  5. The PCIe bottleneck: The 8 GPUs communicate via PCIe (no NVLink), making all-reduce operations extremely expensive. This is the fundamental constraint driving the entire optimization effort — every all-reduce costs ~25ms, and the verify step requires 122 of them.

Output Knowledge Created

This message produces several concrete outputs:

  1. A confirmed hypothesis: The vocab mappings between K2 and K2.5 are indeed substantially different (99.2% mismatch). This was expected but unconfirmed until now.
  2. A validated approach: The decision to use AQ-MedAI's native mapping for Phase 0 is sound. The probe will measure hidden state similarity without the confounding factor of a mismatched vocabulary mapping.
  3. A fixed configuration: The AQ-MedAI drafter's config.json now has max_position_embeddings: 262144, making it compatible with SGLang's context length requirements.
  4. A documented reasoning chain: The assistant's internal monologue — from initial observation, to deeper realization, to correct conclusion — is preserved in the message, providing a template for future debugging of similar issues.

Broader Significance

What makes this message noteworthy is not the complexity of the code change (a simple JSON edit) but the quality of the reasoning. The assistant could have easily panicked at the 99.2% mismatch and concluded that the AQ-MedAI drafter was incompatible with K2.5. Instead, it correctly identified that the mapping should be different because it reflects the training data distribution, not an architectural property. The lm_head learned to predict tokens in the coordinate system defined by AQ-MedAI's mapping, so using that same mapping is correct.

This kind of reasoning — distinguishing between "something is wrong" and "something is different" — is the essence of effective debugging in ML systems. The assistant demonstrated an understanding that neural network weights encode relationships relative to their training context, and that changing the context (the vocabulary mapping) without updating the weights would break those relationships.

The message also illustrates a broader pattern in this session: the constant tension between the desire for quick results (Phase 0 is supposed to take 30 minutes) and the need for careful analysis. The assistant could have skipped the vocab mapping check and launched the server, only to encounter mysterious failures or poor performance. Instead, it invested a few extra minutes in verification, potentially saving hours of debugging.

In the end, the assistant launches the server with the AQ-MedAI drafter (in [msg 4951]), and the probe proceeds. The vocab mapping question was resolved correctly, allowing the experiment to measure what it was supposed to measure: the similarity of hidden state representations between K2 and K2.5, not the compatibility of their vocabulary mappings.