The Config That Bridges Training and Deployment: Making an EAGLE-3 Draft Model vLLM-Compatible

In the course of deploying a 1-trillion-parameter Kimi-K2.5 model on 8x Blackwell GPUs, a critical sub-project emerged: building an EAGLE-3 speculative decoding draft model to accelerate inference. After constructing a complete training pipeline using the speculators library and validating it on 1,000 samples, the assistant arrived at a moment of truth—could the trained checkpoint actually be loaded by vLLM for inference? Message 2794 captures this pivotal transition, where the assistant writes a flat vLLM-compatible configuration file to bridge the gap between the training framework's output format and the inference engine's expectations.

The EAGLE-3 Training Pipeline: A Quick Retrospective

To understand message 2794, one must appreciate the journey that preceded it. The assistant had been working for multiple segments on deploying Kimi-K2.5, a massive MoE (Mixture-of-Experts) model, and had identified the AllReduce communication bottleneck as the primary limiter of decode throughput. Speculative decoding emerged as a promising software-only optimization path, and EAGLE-3—a draft model architecture that predicts multiple future tokens using a lightweight transformer conditioned on the verifier model's hidden states—was chosen as the approach.

The training pipeline consisted of several stages: extracting hidden states from the verifier model (Kimi-K2.5 itself), building vocabulary mappings, and training a small EAGLE-3 draft model using the speculators library. The assistant had rewritten 04_train.py to properly use the speculators API (Eagle3SpeculatorConfig, Eagle3DraftModel, and the built-in Trainer class), monkey-patching the verifier weight extraction to handle Kimi-K2.5's nested configuration structure. The pipeline was validated on 10 samples (3 epochs in ~1 minute), then scaled to 1,000 samples with full hidden state extraction (22.5 minutes for model load, 2.9 minutes for extraction at 2,912 tok/s) and training (10 epochs in 27.7 minutes at 6 steps/s).

The output checkpoint was a 4.6 GB model.safetensors file containing 16 weight tensors, including the fusion layer (fc.weight), attention projections (q_proj, k_proj, v_proj, o_proj), MLP layers, and vocabulary mapping buffers (d2t). The weight shapes had been verified to match the AQ-MedAI reference model exactly—a critical compatibility milestone.

The Config Format Problem

However, a subtle but important issue remained. The speculators library's save_pretrained() method produced a nested config format with a speculators_config object containing algorithm-specific parameters and a transformer_layer_config for the underlying transformer architecture. The architectures field was set to ["Eagle3DraftModel"], which is the speculators library's internal model class name.

vLLM, on the other hand, expected a flat config format with architectures: ["LlamaForCausalLMEagle3"]—the vLLM-specific model class that handles EAGLE-3 speculative decoding. The AQ-MedAI reference model used this flat format, with all transformer parameters (hidden_size, num_attention_heads, head_dim, etc.) at the top level, and an eagle_config nested object containing eagle_aux_hidden_state_layer_ids and use_aux_hidden_state.

The assistant had discovered this discrepancy in earlier messages ([msg 2786] and [msg 2787]), where it compared the speculators output format against the AQ-MedAI reference. While vLLM did have a SpeculatorsConfig handler that could theoretically read the nested format, the flat format was the proven path—it was what the reference model used, and it was what vLLM's LlamaForCausalLMEagle3 class expected.

The Subject Message: Writing the Flat Config

Message 2794 is the direct response to this discovered incompatibility. The assistant executes a single bash command over SSH that runs an inline Python script on the remote machine. The script performs two operations:

Backs up the speculators config: It copies the existing config.json (the nested speculators format) to config_speculators.json, preserving the original for reference or fallback.
Writes a flat vLLM-compatible config: It constructs a complete configuration dictionary that mirrors the AQ-MedAI format exactly, with all the parameters vLLM needs to instantiate the LlamaForCausalLMEagle3 model. The config dictionary contains 30 fields, each carefully chosen: - Architecture identification: architectures: ["LlamaForCausalLMEagle3"] tells vLLM which model class to use. - EAGLE-specific parameters: eagle_config.eagle_aux_hidden_state_layer_ids: [2, 30, 58] specifies which layers of the verifier model provide hidden states to the draft model. These layer IDs (2, 30, 58) correspond to the early, middle, and late layers of Kimi-K2.5's 60-layer architecture, providing diverse representational inputs. - Transformer architecture: hidden_size: 7168, num_attention_heads: 64, num_key_value_heads: 64, head_dim: 128, intermediate_size: 18432, num_hidden_layers: 1 define the draft model's single-layer transformer. - Vocabulary: vocab_size: 163840 (Kimi-K2.5's full vocabulary) and draft_vocab_size: 32000 (the draft model's smaller vocabulary) define the token spaces. - RoPE configuration: rope_theta: 1000000.0 and max_position_embeddings: 131072 match the base model's long-context capabilities. - Miscellaneous: attention_bias: False, mlp_bias: False, rms_norm_eps: 1e-6, tie_word_embeddings: False—standard settings that must match the training configuration. The head_dim: 128 field is particularly noteworthy. In a standard Llama configuration, head_dim is computed as hidden_size / num_attention_heads = 7168 / 64 = 112. But the AQ-MedAI reference explicitly sets it to 128, meaning the Q/K/V projections project up from 7168 to 8192 (64 heads × 128). The assistant had discovered this in [msg 2781] and corrected the training script accordingly, ensuring the weight shapes matched the reference exactly.

Assumptions and Knowledge

This message embodies several key assumptions:

Assumption 1: The flat format is necessary for vLLM compatibility. While vLLM does have a SpeculatorsConfig handler that can read the nested format, the assistant chose the flat format for maximum reliability. This was a conservative, safe choice—mirroring the proven AQ-MedAI format rather than relying on less-tested code paths.

Assumption 2: The config values must exactly match the training configuration. Any discrepancy between the config used during training and the config used during inference could cause silent failures—wrong weight shapes, incorrect initialization, or runtime errors. The assistant carefully replicated every parameter from the training setup.

Assumption 3: Backing up the original config is sufficient. The assistant preserves the speculators config as config_speculators.json, but the primary config.json is overwritten. This assumes that the flat format is the definitive one and that the speculators format is only needed for reference.

Input knowledge required to understand this message includes: familiarity with the EAGLE-3 architecture (a single transformer layer conditioned on verifier hidden states), knowledge of vLLM's speculative decoding infrastructure, understanding of the speculators library's training API and output format, and awareness of the AQ-MedAI reference model as the canonical format to replicate.

Output knowledge created by this message is a vLLM-compatible checkpoint directory at /root/eagle3-train/output_test/2/ that can be loaded by vLLM's LlamaForCausalLMEagle3 class. This checkpoint is the bridge between the training pipeline and the inference deployment—without it, the trained draft model would be unusable.

The Thinking Process

The assistant's reasoning is visible in the sequence of actions leading to this message. In [msg 2785], a task tool investigated how vLLM loads EAGLE-3 models, discovering that vLLM accepts both midlayer.* and layers.0.* weight naming conventions and that the embed_tokens weight is optional. In [msg 2786], the assistant inspected the AQ-MedAI config directly, comparing it field-by-field with the speculators output. In [msg 2787], the assistant explicitly noted the differences and considered whether to produce the flat format.

The decision to write the flat config was not made lightly. The assistant could have relied on vLLM's SpeculatorsConfig handler, which would read the nested format. But the flat format had several advantages: it was the proven path (AQ-MedAI used it), it avoided any potential bugs in the less-tested speculators config handler, and it made the checkpoint directly compatible with any vLLM version without requiring speculators library support.

The message itself is concise—a single bash command—but it represents the culmination of hours of investigation, debugging, and validation. The config dictionary is not copied from AQ-MedAI blindly; each field has been verified against the training configuration and the reference model. The eagle_aux_hidden_state_layer_ids: [2, 30, 58] matches Kimi-K2.5's specific architecture. The head_dim: 128 corrects the default computation. The vocab_size: 163840 matches Kimi-K2.5's extended vocabulary.

Significance in the Larger Context

Message 2794 is a turning point in the EAGLE-3 sub-project. Up to this point, the work had been entirely in the training domain—extracting hidden states, configuring the speculators library, running training loops. This message marks the transition to the deployment domain, where the trained checkpoint must integrate with the inference engine.

The next logical step, which the assistant explicitly contemplated in [msg 2793], is to actually test loading the checkpoint with vLLM and serving it in speculative decoding mode. This would validate the entire pipeline end-to-end: hidden state extraction → training → config conversion → inference. Without message 2794, that validation would fail on a config format mismatch, wasting the hours invested in training.

The config format bridge is a small piece of engineering—a single JSON file—but it is the kind of detail that determines whether a complex pipeline actually works end-to-end. The assistant's careful attention to this compatibility layer, informed by direct inspection of both the speculators output and the AQ-MedAI reference, exemplifies the thoroughness required to deploy cutting-edge ML systems in production.