The Weight-Shape Verdict: Validating EAGLE-3 Training Against the AQ-MedAI Reference

In the long arc of building a speculative decoding pipeline for the Kimi-K2.5 model on 8× Blackwell GPUs, there comes a moment of quiet triumph that is easy to miss amid the torrent of bash commands, SSH sessions, and config edits. Message 2791 is that moment. It is a brief, almost clinical message from the assistant, yet it represents the culmination of dozens of preceding steps: exploring the speculators library API, rewriting training scripts, monkey-patching config extraction, running hidden-state generation across thousands of samples, and debugging subtle architectural mismatches. This message is the validation gate—the point where the assistant confirms that its EAGLE-3 training pipeline produces output that is structurally identical to the reference implementation from AQ-MedAI, and then takes the final step to ensure maximum compatibility with the vLLM inference engine.

The Context: Building a Custom EAGLE-3 Draft Model

To understand why message 2791 matters, one must appreciate the journey that led to it. The assistant had been tasked with deploying and optimizing large language models on a machine with 8× NVIDIA RTX PRO 6000 Blackwell GPUs. After deploying the Kimi-K2.5 model in its INT4 quantized variant and profiling its performance, the dominant bottleneck was identified: AllReduce operations consumed 51.5% of decode time. Speculative decoding—using a smaller "draft" model to propose tokens that a larger "verifier" model then accepts or rejects—emerged as the most promising software-only optimization path.

The assistant chose EAGLE-3, a state-of-the-art speculative decoding architecture that uses a lightweight transformer module to predict hidden states from the verifier's own intermediate layers. But implementing EAGLE-3 for Kimi-K2.5 required building a complete training pipeline from scratch, adapting the speculators library (v0.3.0) to work with vLLM 0.16, and resolving numerous API incompatibilities along the way. By message 2791, the assistant had already:

1. Explored the speculators library's __main__.py and Trainer class to understand the proper training workflow
2. Rewritten 04_train.py to use Eagle3SpeculatorConfig, Eagle3DraftModel, and the built-in Trainer
3. Validated training on 10 samples (3 epochs in ~1 minute)
4. Scaled to 1000 samples with full hidden state extraction (22.5 minutes for model load, 2.9 minutes for extraction at 2912 tok/s, and 27.7 minutes for training) But a critical question remained: would the trained checkpoint actually work with vLLM? The speculators library and vLLM are developed by different teams, and there is no guarantee that a checkpoint saved by one can be loaded by the other.

The Verification: Matching Weight Shapes

Message 2791 opens with the assistant's conclusion after running a training iteration with head_dim=128:

Now the weight shapes exactly match AQ-MedAI's reference:

The assistant then presents a comparison table showing six key weight tensors from its own checkpoint and from the AQ-MedAI reference model (AQ-MedAI/Kimi-K2-Instruct-eagle3). Every shape is identical:

• fc.weight: [7168, 21504] — the fusion layer that combines the verifier's hidden states with the draft model's embeddings
• q_proj.weight: [8192, 14336] — the query projection, projecting from 2× hidden_size (14336) up to num_heads × head_dim (8192)
• k_proj.weight: [8192, 14336] — same dimensions for the key projection
• v_proj.weight: [8192, 14336] — same for the value projection
• o_proj.weight: [7168, 8192] — the output projection, mapping back to hidden_size
• lm_head.weight: [32000, 7168] — the language modeling head mapping to the draft vocabulary This exact match is not trivial. It required the assistant to notice and correct a subtle architectural discrepancy earlier in the session. In message 2780, the assistant had observed that the initial training run used head_dim=112 (computed implicitly as hidden_size / num_attention_heads = 7168 / 64), but the AQ-MedAI reference explicitly sets head_dim=128. This means the attention projections are wider than the default—projecting from 7168 up to 8192 dimensions rather than down to 7168. The assistant recognized this as an intentional architectural choice in the reference model and corrected the config accordingly (message 2781), then re-ran the training (messages 2788-2790).

The Remaining Differences and Their Significance

After presenting the shape comparison, the assistant notes two remaining differences:

The only differences are naming convention (we use layers.0.*, they use midlayer.* — both work in vLLM) and we include embed_tokens.weight (optional, vLLM shares from verifier if missing or identical).

These differences are worth examining because they reveal the assistant's deep understanding of the vLLM loading pipeline. The naming convention difference (layers.0.* vs midlayer.*) stems from how the speculators library names its transformer layers versus how AQ-MedAI's custom implementation names them. The assistant had already verified in a subagent task (message 2784) that vLLM's Eagle3Loader explicitly remaps midlayer to layers.0 during loading, so both naming conventions work. The embed_tokens.weight inclusion is similarly benign—vLLM will share the verifier's embedding table if the draft checkpoint doesn't provide one, or use the provided one if it matches.

The Decision: Producing a Flat Config

The most consequential part of message 2791 is the assistant's decision to add a conversion step:

Now let me also produce a vLLM-compatible flat config alongside the speculators config, so we have maximum compatibility. Let me create a conversion step that produces the AQ-MedAI-style flat config:

This decision was informed by the assistant's earlier investigation of the AQ-MedAI config format (message 2786). The reference model uses a flat config structure where all model hyperparameters (hidden_size, num_attention_heads, head_dim, etc.) sit at the top level of the JSON, with a nested eagle_config object containing the EAGLE-specific parameters (eagle_aux_hidden_state_layer_ids, use_aux_hidden_state). The architecture is declared as "LlamaForCausalLMEagle3"—a type that vLLM recognizes directly.

In contrast, the speculators library's save_pretrained() produces a nested config format with speculators_config and transformer_layer_config sub-objects, and declares the architecture as "Eagle3DraftModel". While vLLM does have a SpeculatorsConfig handler that can parse this nested format, the assistant recognized that the flat format provides maximum compatibility—it matches exactly what the reference model uses, and it avoids any risk of vLLM misinterpreting the nested structure.

The assistant's reasoning here reveals a pragmatic engineering mindset: rather than assuming that the speculators format will work (even though it should, based on the code inspection), the assistant chooses to produce both formats. This dual-format approach hedges against the possibility that vLLM's speculators config handler has edge cases or bugs that haven't been encountered yet.

Assumptions and Knowledge

Message 2791 rests on several key assumptions and bodies of knowledge:

Input knowledge required:

• Understanding of the EAGLE-3 architecture: the fusion layer (fc.weight), the attention projections (Q/K/V/O), the language modeling head, and how they map to specific tensor shapes
• Familiarity with the AQ-MedAI reference model and its config format
• Knowledge of vLLM's speculative decoding infrastructure, particularly the Eagle3Loader class and its weight remapping logic
• Understanding of the speculators library's save_pretrained() output format
• Awareness of the head_dim parameter and how it affects attention projection dimensions Assumptions made:
• That vLLM's handling of both naming conventions (layers.0.* and midlayer.*) is correct and will not change between versions
• That producing a flat config alongside the speculators config is sufficient for compatibility—i.e., that vLLM will prefer one format over the other
• That the weight shape match is sufficient to guarantee functional correctness (the assistant does not run an actual inference test with the trained checkpoint in this message)
• That the LSP errors shown in the diagnostics are false positives caused by the remote development environment not having the Python packages installed locally Output knowledge created:
• Confirmation that the training pipeline produces checkpoints structurally identical to the reference
• A modified training script that produces both config formats
• Documentation (via the comparison table) of the exact weight shapes expected by vLLM for EAGLE-3 on Kimi-K2.5

The Thinking Process

The assistant's reasoning in message 2791 follows a clear pattern: verify, compare, decide, act. First, it verifies the weight shapes from the re-run training (message 2790). Then it compares them against the AQ-MedAI reference that was downloaded and inspected earlier (message 2783). Based on the comparison, it decides that the weight shapes are correct but the config format could be improved. Finally, it acts by editing the training script to add the flat config conversion.

What is notable about this message is what it does not contain. There is no bash command, no SSH session, no output to inspect. The assistant has already done the verification work in the preceding messages. Message 2791 is purely analytical and decision-making—it synthesizes the results of previous investigations and takes a proactive step to improve compatibility.

The comparison table is particularly well-constructed. By presenting the shapes side-by-side, the assistant makes the match visually unambiguous. This is a communication choice that serves multiple purposes: it documents the expected shapes for future reference, it provides clear evidence that the training pipeline is correct, and it implicitly justifies the decision to proceed with the flat config conversion.

The Broader Significance

In the context of the entire coding session, message 2791 marks the transition from building the training pipeline to preparing for inference integration. Up to this point, the assistant's focus has been on getting the training to work—resolving API incompatibilities, fixing config mismatches, and validating the training loop. After this message, the assistant will pivot to generating high-quality synthetic training data (as noted in the chunk summary) and eventually deploying the trained EAGLE-3 draft model for speculative decoding.

The message also illustrates a broader principle of ML engineering: checkpoint compatibility is not guaranteed. When working with models that span multiple frameworks (speculators for training, vLLM for inference), the burden falls on the engineer to verify that the output of one system can be consumed by the other. The assistant's systematic approach—downloading the reference checkpoint, comparing weight shapes, inspecting config formats, and testing both naming conventions—is a model of how to approach cross-framework compatibility.

Finally, message 2791 demonstrates the value of reference implementations. Without the AQ-MedAI checkpoint to compare against, the assistant would have had to guess whether the weight shapes were correct. The reference provided an objective standard, and the assistant leveraged it thoroughly—not just for the shape comparison, but also for the config format and the architectural parameters like head_dim.

Conclusion

Message 2791 is a validation milestone in a complex ML engineering effort. It confirms that the EAGLE-3 training pipeline produces checkpoints that are structurally identical to the reference implementation, and it takes a proactive step to ensure maximum compatibility with the vLLM inference engine. The message is concise but dense with meaning: it encapsulates hours of debugging, multiple re-runs of training, and deep investigation of both the speculators library and vLLM's loading code. For anyone following the session, it provides the confidence that the training pipeline is correct and that the next phase—generating high-quality training data and deploying the draft model—can proceed on solid ground.