Bridging Training and Inference: Verifying EAGLE-3 Checkpoint Compatibility with vLLM

In the complex pipeline of building a speculative decoding system for a 1-trillion-parameter model like Kimi-K2.5, one of the most critical junctures occurs when the training pipeline hands off its output to the inference engine. Message 2785 captures exactly this moment: the assistant has just completed a successful end-to-end training run of an EAGLE-3 draft model on 10 samples, and is now meticulously verifying that the resulting checkpoint is compatible with vLLM's loading code. This message is a synthesis of findings from a subagent task, a compatibility analysis, and a forward-looking plan to fix a subtle but critical configuration mismatch.

The Message in Full

The message opens with the assistant summarizing six key findings about how vLLM loads EAGLE-3 speculator models, drawn from a subagent task that explored the vLLM codebase:

This is great news. Let me summarize the key findings for compatibility:

>

1. Weight naming: vLLM accepts both midlayer.* AND layers.0.* (it remaps midlayer to layers.0). Our layers.0.* naming works. 2. embed_tokens: Optional — if missing from checkpoint, vLLM shares from verifier. We can save it or not. 3. lm_head: The draft lm_head (draft_vocab_size) is loaded if present. 4. d2t: Renamed to draft_id_to_target_id internally, used for vocab mapping. 5. t2d: Completely skipped/ignored. 6. Config: vLLM reads eagle_aux_hidden_state_layer_ids from the draft config.

>

The main issue is that our test run used head_dim=112 (implicit from 7168/64). I need to re-run with head_dim=128 to match AQ-MedAI.

The assistant then executes a bash command to inspect the saved config.json, and the output reveals the full configuration structure including the speculators_config block with algorithm settings, proposal methods, and acceptance tolerances.

Why This Message Was Written

This message exists at a specific inflection point in the development workflow. The assistant had just completed a grueling multi-step process: exploring the speculators library's training API, rewriting 04_train.py to use proper Eagle3SpeculatorConfig and Trainer classes, monkey-patching verifier weight extraction for Kimi-K2.5's nested config structure, and successfully running a 3-epoch training loop on 10 samples in about one minute. The training pipeline was validated — it produced a 4.6 GB checkpoint with the correct weight shapes and a proper HuggingFace-compatible format.

But training success is only half the battle. The checkpoint must be consumable by vLLM, the inference engine that will actually use the EAGLE-3 draft model for speculative decoding at runtime. If the checkpoint format, weight naming, or configuration structure doesn't match what vLLM expects, the entire training effort is wasted. The assistant recognized this and dispatched a subagent task to investigate vLLM's EAGLE-3 loading code. Message 2785 is the synthesis of those findings and the resulting action plan.

The urgency is palpable in the message's structure. The assistant doesn't pause to celebrate the training success — it immediately pivots to compatibility verification, treating the training run as merely a stepping stone toward the real goal of deploying speculative decoding. This reflects a deep understanding of the end-to-end system: in production ML engineering, a model that trains perfectly but cannot be loaded for inference is worthless.

The Thinking Process: Systematic Compatibility Analysis

The assistant's reasoning in this message reveals a methodical, checklist-driven approach to compatibility verification. Each finding addresses a distinct potential failure mode:

Finding 1 (Weight naming): The assistant had noticed earlier that AQ-MedAI's reference checkpoint uses midlayer.* prefixes while the speculators library's save_pretrained() produces layers.0.*. This could have been a showstopper — if vLLM hard-codes midlayer lookups, all weights would fail to load. The subagent confirmed that vLLM's Eagle3DraftModel.load_weights() method remaps midlayer to layers.0, meaning both naming conventions work. This is a relief, but the assistant files it away as verified rather than assumed.

Finding 2 (embed_tokens): The embed_tokens.weight tensor is a 163,840×7,168 matrix — a massive embedding table that consumes significant storage. If vLLM requires it, the checkpoint must include it. If vLLM can share it from the verifier model, the assistant has the option to omit it and save disk space. The subagent confirmed it's optional, giving the assistant flexibility.

Finding 3 (lm_head): The draft model's language modeling head maps hidden states to vocabulary logits. Its presence or absence affects how vLLM initializes the draft model. The subagent confirmed it's loaded if present, meaning the assistant can include it for correctness without worrying about conflicts.

Finding 4 (d2t → draft_id_to_target_id): The d2t buffer maps draft vocabulary IDs to target vocabulary IDs — this is the core of EAGLE-3's vocabulary projection mechanism. The assistant had saved it as d2t in the checkpoint, and the subagent confirmed vLLM renames it to draft_id_to_target_id internally. The shape [32000] matches the draft vocabulary size, confirming correctness.

Finding 5 (t2d): The t2d buffer (target-to-draft mapping) is completely ignored by vLLM. The assistant notes this without judgment — it's simply information that can guide future checkpoint optimization.

Finding 6 (Config — eagle_aux_hidden_state_layer_ids): This is perhaps the most critical config field. It tells the draft model which layers of the verifier to extract hidden states from. For Kimi-K2.5, the assistant had set this to layers 2, 30, and 58 — the 2nd, 30th, and 58th transformer layers. If vLLM doesn't read this field, the draft model would extract states from the wrong layers and produce garbage predictions. The subagent confirmed vLLM does read it, validating the training configuration.

The Critical Head_Dim Discrepancy

The most important finding is the head_dim mismatch. The assistant's initial training run used the mathematically natural value: head_dim = hidden_size / num_attention_heads = 7168 / 64 = 112. However, the AQ-MedAI reference model uses head_dim = 128, which means the attention projections are 64×128 = 8192 dimensions — larger than the hidden size of 7168. This is an intentional design choice that makes the Q/K/V projections project up from 7168 to 8192, giving the attention mechanism more representational capacity.

The assistant recognizes this as "the main issue" and immediately plans to fix it. This is a subtle but critical detail: if the head_dim doesn't match, the weight shapes of q_proj, k_proj, v_proj, and o_proj will all be wrong, and vLLM will either fail to load the checkpoint or produce incorrect attention computations. The assistant had already edited the training script to set head_dim=128 in a previous message (msg 2781), but the test run predated that edit. Now the assistant must re-run training with the corrected configuration.

Assumptions and Potential Pitfalls

Several assumptions underpin the assistant's analysis. First, the assistant assumes that matching AQ-MedAI's format is the correct target. This is reasonable — AQ-MedAI's model is the only publicly available EAGLE-3 checkpoint, and vLLM's loading code was likely developed against it. However, it's possible that vLLM's loading code is flexible enough to handle different configurations, and the assistant's checkpoint might have worked even with head_dim=112. The assistant doesn't test this assumption — it defaults to matching the reference.

Second, the assistant assumes that the vLLM loading code it examined is the final, authoritative implementation. The subagent explored a specific version of vLLM (0.16.x, installed in the ML environment). If the vLLM version changes, the loading code might change too. The assistant doesn't document which version was examined, which could cause confusion if compatibility issues arise later with a different vLLM version.

Third, the assistant assumes that the config.json format produced by save_pretrained() is what vLLM expects. The output shows a speculators_config block with algorithm settings, proposal methods, and acceptance tolerances. The assistant doesn't verify that vLLM reads these fields — it only confirms that eagle_aux_hidden_state_layer_ids is read. Other fields like norm_before_residual and has_no_defaults_at_init might be silently ignored or might cause loading failures.

Input and Output Knowledge

Input knowledge required to understand this message includes:

The Broader Context

This message sits within a larger narrative of building a production-grade speculative decoding system. The assistant had previously resolved API incompatibilities between the speculators library and vLLM 0.16, patched the custom worker for Kimi-K2.5's DeepSeekV2 architecture, and successfully extracted hidden states from 10 test samples. The training pipeline was validated end-to-end, and now the assistant is ensuring the pipeline's output integrates seamlessly with the inference stack.

The config.json output shown in the message reveals the full complexity of the EAGLE-3 configuration: a speculators_config block with algorithm type "eagle3", a default proposal method of "greedy", an acceptance tolerance of 0.0 (meaning only exact token matches are accepted), and a draft vocabulary of 32,000 tokens mapped from the target vocabulary. The eagle_aux_hidden_state_layer_ids of [2, 30, 58] reflects a deliberate choice to extract features from early, middle, and late layers of the 61-layer Kimi-K2.5 model, giving the draft model multi-scale representations of the verifier's reasoning process.

Conclusion

Message 2785 exemplifies the meticulous, systematic approach required when building complex ML systems. The assistant doesn't assume compatibility — it verifies each dimension of the checkpoint format against the inference engine's expectations, identifies a subtle but critical configuration mismatch, and plans the correction. This message is a bridge between two worlds: the training pipeline that produces the draft model, and the inference engine that consumes it. By carefully documenting the compatibility findings and the head_dim issue, the assistant ensures that the next training run will produce a checkpoint that vLLM can load and use for speculative decoding, bringing the project one step closer to its goal of accelerating Kimi-K2.5 inference on 8 Blackwell GPUs.