The Vocabulary Bridge: Building an EAGLE-3 Draft Model for Kimi-K2.5

Introduction

In the high-stakes world of large language model inference optimization, every millisecond counts. When profiling revealed that AllReduce communication consumed 51.5% of decode time on an 8×RTX PRO 6000 Blackwell system running Kimi-K2.5 INT4, the natural next step was speculative decoding — a technique where a small "draft" model proposes tokens that a large "target" model verifies in parallel, trading communication-bound decode for compute-bound verification. But for a 1-trillion-parameter MoE model like Kimi-K2.5, no off-the-shelf draft model exists. The only viable path was to train a custom EAGLE-3 draft head, and that meant building a complete training pipeline from scratch.

Message [msg 2532] captures a quiet but crucial moment in that pipeline construction: the successful execution of two non-GPU pipeline steps — dataset preparation and vocabulary mapping — that together validate the data foundation upon which the entire EAGLE-3 training effort rests. This message is a checkpoint of progress, a validation of assumptions, and a bridge between the abstract plan and the concrete implementation.

The Message in Full

The assistant reports:

Step 1 works. 10 samples processed, 3875 tokens total, 2875 assistant tokens. The tokenizer's custom encode is chatty with debug logs but works fine.

>

Now Step 3 (vocab mapping — doesn't need GPU):

>

``` [bash] ssh root@10.1.230.174 '~/ml-env/bin/python3 /root/eagle3-train/03_build_vocab_mapping.py --token-freq /root/eagle3-train/data_test/prepared/token_freq.pt --output-dir /root/eagle3-train/data_test/vocab_mapping --target-vocab-size 163840 --draft-vocab-size 32000 2>&1' Loading token frequency from /root/eagle3-train/data_test/prepared/token_freq.pt... Shape: torch.Size([163840]) Non-zero tokens: 899 Total frequency: 2875

>

Building vocab mapping: 163840 -> 32000... Draft vocab covers 100.0% of token frequency Selected token ID range: [0, 163607] t2d shape: torch.Size([163840]), dtype: torch.bool d2t shape: torch.Size([32000]), dtype: torch.int64

>

Saved to /root/eagle3-train/data_test/vocab_mapping/ t2d.pt, d2t.pt (torch), t2d.npy, d2t.npy (numpy... ```

Two commands, two confirmations. On the surface, it reads like a routine progress update. But beneath the terseness lies a dense web of technical decisions, architectural assumptions, and carefully managed risk.

Context and Motivation: Why This Message Exists

To understand why [msg 2532] was written, we must trace back through the preceding messages. The assistant had just completed a deep profiling campaign of Kimi-K2.5 INT4 on 8×Blackwell GPUs ([msg 2507] and surrounding context in Segment 19), which identified AllReduce as the dominant bottleneck at 51.5% of decode time. The user then pivoted to investigate speculative decoding as a software-only optimization path.

The research phase ([msg 2508] through [msg 2516]) was thorough: the assistant launched parallel research agents covering speculative decoding fundamentals, vLLM/SGLang framework support, candidate draft models, and training feasibility. The key findings were sobering:

The Vocabulary Mapping Problem: Why 163,840 → 32,000?

The core technical challenge addressed in this message is the vocabulary mismatch between the target model (Kimi-K2.5) and the draft model (EAGLE-3). Kimi-K2.5 uses a vocabulary of 163,840 tokens — a large vocabulary typical of modern multilingual models trained on diverse data. The EAGLE-3 draft model, by contrast, uses a much smaller vocabulary of 32,000 tokens.

This mismatch creates a fundamental problem: when the draft model proposes a token, it proposes a token ID from its own 32K vocabulary. But the target model needs to verify that token using its 164K vocabulary. The two vocabularies overlap significantly — the draft vocabulary is essentially a subset of the most frequent tokens from the target vocabulary — but the mapping is not identity. Token ID 42 in the draft vocabulary might correspond to token ID 1423 in the target vocabulary, or it might not exist at all.

The vocabulary mapping script (03_build_vocab_mapping.py) solves this by:

  1. Loading the token frequency distribution from the prepared dataset
  2. Building a boolean "target-to-draft" (t2d) tensor of shape [163840] that marks which target tokens are in the draft vocabulary
  3. Building an integer "draft-to-target" (d2t) tensor of shape [32000] that maps each draft token ID to its corresponding target token ID The output confirms that the draft vocabulary covers 100.0% of token frequency — meaning every token that appeared in the 10-sample test dataset is representable in the draft vocabulary. This is expected for a small sample (only 899 unique tokens out of 163,840), but it's a necessary validation that the mapping logic works correctly.

The Token Frequency Distribution: Reading Between the Numbers

The output reveals interesting statistics about the test data:

Assumptions and Design Decisions

Several implicit assumptions underpin this message:

1. The EAGLE-3 architecture assumption. The assistant assumes that an EAGLE-3 draft head with a 32K vocabulary is the right architecture for Kimi-K2.5. This is based on the existing AQ-MedAI/Kimi-K2-Instruct-eagle3 model, which uses the same architecture for the similar K2 model. The assumption is that K2.5, being an incremental improvement over K2, would benefit from the same draft model design. This is reasonable but unverified — the optimal draft vocabulary size for K2.5 might differ.

2. The token frequency heuristic. The vocabulary mapping uses token frequency from the training dataset to determine which target tokens map to which draft tokens. This assumes that the most frequent tokens in the training data are the most important ones for the draft model to know. While this is standard practice in speculative decoding, it's a heuristic — the draft model might benefit more from knowing rare but semantically important tokens.

3. The small-sample validation strategy. The assistant tests with only 10 samples from mlabonne/open-perfectblend, a dataset of 2,500 perfect blend examples. This is a deliberate risk management strategy: validate the pipeline logic with minimal data before scaling up. The assumption is that pipeline correctness generalizes from 10 samples to thousands. This is reasonable for data processing pipelines but less so for training — the hidden state extraction and training steps may reveal issues that don't surface with tiny datasets.

4. The offline extraction approach. The assistant chose to write a custom pipeline using speculators' VllmHiddenStatesGenerator rather than trying to extract hidden states from the running vLLM server. This decision was made in [msg 2516] after discovering that speculators' data generation code monkey-patches vLLM internals and would be incompatible with vLLM 0.16. The assumption is that stopping the server, running extraction, and restarting is acceptable downtime. For a production deployment this would be problematic, but for a research/training pipeline it's pragmatic.

What Input Knowledge Is Required

To fully understand this message, one needs:

  1. EAGLE-3 architecture knowledge. Understanding that EAGLE-3 is a speculative decoding framework where a lightweight draft model predicts tokens using hidden states from the target model's intermediate layers. The draft model has its own vocabulary (typically smaller than the target's) and needs a mapping between the two.
  2. Vocabulary mapping mechanics. Knowing that large language models use tokenizers with fixed vocabularies, and that when two models have different vocabularies, a mapping must be constructed. The mapping is typically based on token frequency — the most common target tokens are assigned to draft vocabulary slots.
  3. The Kimi-K2.5 model context. Understanding that Kimi-K2.5 is a 1T-parameter MoE model from Moonshot AI, that it uses a 163,840-token vocabulary, and that it's deployed with INT4 quantization on 8×Blackwell GPUs. The assistant has been working with this model throughout the session.
  4. The pipeline architecture. Knowing that the EAGLE-3 training pipeline consists of five stages: draft config creation, dataset preparation, hidden state extraction, vocabulary mapping, and training. Steps 1 and 3 (the ones tested here) are the data preparation stages that don't require GPU.
  5. The speculators library. Understanding that speculators is a vLLM-adjacent library for speculative decoding training, and that its VllmHiddenStatesGenerator class handles the complex task of extracting intermediate hidden states from a running vLLM instance.

What Output Knowledge Is Created

This message produces several concrete artifacts:

  1. Validation of Step 1 (dataset preparation). The assistant confirms that the dataset preparation script correctly processes 10 samples, tokenizes them, separates prompt from assistant tokens, and produces the expected token frequency distribution. The "chatty" tokenizer debug output is noted but deemed non-blocking.
  2. Vocabulary mapping tensors. The script produces four files: - t2d.pt / t2d.npy: A boolean tensor of shape [163840] marking which target tokens are in the draft vocabulary - d2t.pt / d2t.npy: An integer tensor of shape [32000] mapping each draft token to its target equivalent Both PyTorch and NumPy formats are saved for flexibility.
  3. Coverage confirmation. The draft vocabulary covers 100% of token frequency in the test dataset, confirming that the mapping logic works and that the 32K draft vocabulary is sufficient for the sampled data.
  4. A verified pipeline checkpoint. The assistant now knows that Steps 1 and 3 work correctly, and can proceed to the GPU-bound Step 2 (hidden state extraction) with confidence that the data foundation is solid.
  5. Documentation of the approach. The commands and outputs in this message serve as a record of the pipeline's behavior, useful for debugging later or for reproducing the setup on different hardware.

Mistakes and Incorrect Assumptions

While the message itself shows successful execution, several potential issues lurk beneath the surface:

1. The 100% coverage figure is misleading. The draft vocabulary covers 100% of token frequency in the test dataset, but this is because the test dataset is tiny (899 unique tokens out of 163,840). With a full training dataset, coverage would be significantly lower — many rare tokens would fall outside the 32K draft vocabulary. The assistant doesn't acknowledge this limitation explicitly, though it's likely understood given the context of a "test" run.

2. The token frequency distribution may be skewed. The test dataset (mlabonne/open-perfectblend) consists of 2,500 "perfect blend" examples — likely curated, high-quality conversations. The token frequency distribution from this dataset may not match the distribution of tokens the model will encounter during actual inference. If the draft model is trained on this distribution, it may have poor acceptance rates on out-of-distribution inputs.

3. The vocabulary mapping is static. The mapping is computed once from the training data and frozen. During inference, if the model encounters a token not in the draft vocabulary, the draft model cannot propose it — it must propose a substitute. This is inherent to the EAGLE-3 approach, but the assistant doesn't discuss strategies for handling out-of-vocabulary tokens (e.g., fallback to the target model's logits).

4. The "chatty" tokenizer may indicate a configuration issue. The debug output "Calling super().encode with {'add_special_tokens': False}" repeating for every sample suggests that the tokenizer's encode method is printing debug information. While the assistant notes this and says "works fine," excessive debug output could mask real issues or slow down processing at scale. It may also indicate that the tokenizer configuration is non-standard — perhaps the model uses a custom tokenizer that doesn't fully integrate with the HuggingFace transformers library.

The Thinking Process: Strategic Sequencing

The most revealing aspect of [msg 2532] is what it doesn't say. The assistant runs Step 3 immediately after confirming Step 1, deliberately skipping Step 2 (hidden state extraction). This is a strategic sequencing decision: Steps 1 and 3 are CPU-only operations that can be validated quickly and cheaply. Step 2 requires GPU access, which means stopping the running vLLM server, waiting for model loading (~25-30 minutes), and risking GPU memory issues.

By validating the non-GPU steps first, the assistant ensures that when the expensive GPU step runs, it won't fail due to data format issues. This is classic "fail fast, fail cheap" engineering — catch bugs at the earliest possible stage.

The assistant also chose to run the vocabulary mapping before the hidden state extraction, even though logically the mapping depends on the token frequency data (produced by Step 1) and not on the hidden states (produced by Step 2). This ordering is correct: the vocabulary mapping is needed during training to translate between draft and target token IDs, but it doesn't depend on the hidden states themselves.

Conclusion

Message [msg 2532] is a quiet checkpoint in a complex engineering effort. On its surface, it reports the successful execution of two pipeline steps. But beneath that surface lies a carefully orchestrated sequence of technical decisions: the choice of EAGLE-3 architecture, the vocabulary mapping strategy, the small-sample validation approach, and the strategic sequencing of GPU-free validation before GPU-bound execution.

The vocabulary mapping — reducing 163,840 tokens to 32,000 while preserving 100% of observed token frequency — is the critical bridge between the target model's rich vocabulary and the draft model's constrained one. It's a compression step that must be lossless for the observed data and as lossless as possible for unseen data. The 100% coverage figure is both a validation and a warning: it confirms the mapping logic works, but it also hints at the challenges that await with larger, more diverse datasets.

This message exemplifies the kind of methodical, risk-aware engineering that characterizes successful ML infrastructure work. Each step is validated before the next begins. Assumptions are tested cheaply before expensive resources are committed. And the output of each stage is carefully inspected for correctness before being passed to the next. It's not glamorous work — but it's the kind of work that separates successful speculative decoding deployments from failed experiments.