[khazic]

What

Adds EAGLE-3 draft training for a frozen Kimi K3 target, the EAGLE counterpart to the DSpark support in #3263.

Re-opens #3264, which GitHub closed automatically when its base branch huiyingl/feat/kimi-k3 was merged and deleted. Same work, now re-targeted at main; the review history lives on #3264.

Kimi K3 interleaves KDA layers with NoPE MLA full-attention layers, so neither the dense EAGLE-3 draft nor the DeepSeek MLA draft (which always rotates the rope slice) can represent it:

	DeepSeek MLA draft	Kimi K3 draft
rope slice	rotated (interleaved RoPE), per-TTT-step phase offset	carried unrotated, no rotary cache at all
attention output	o_proj(attn)	sigmoid output gate, then o_proj
MLP / norm	SwiGLU, shared RMSNorm helper	K3 SituAndMul MLP, K3 fp32-variance RMSNorm

KimiK3Eagle3DraftModel mirrors exactly one K3 MLA layer (Q-LoRA, the compressed KV latent, the gate, the SiTU MLP); KDA, the 896 routed experts and the attention-residual mixer stay in the target. Everything EAGLE-3 specific (the [embed, hidden] fused first layer, the cache_hidden TTT recurrence with diagonal extension, d2t/t2d, project_hidden_states/compute_logits) is inherited from the DeepSeek MLA draft through one new _build_inner_model hook, so the trainer, the checkpointer and the serve-side contract are reused unchanged.

Positional information reaches the draft through the target hidden states it is conditioned on, exactly as it reaches the target's own MLA layers; test_forward_is_position_id_invariant pins that contract.

Recipe wiring

K3 has no HuggingFace implementation and its experts do not fit on one GPU, so the colocated target is built from the (possibly nested) text config through AutoModel's custom expert-parallel path rather than the generic from_pretrained. A vision-language checkpoint is loaded as its text-only KimiK3ForCausalLM; the target config is a copy, so the draft config serialized into config.json is untouched.

Combinations K3 cannot serve are rejected up front, before the 93-layer target is loaded: non-colocated backends (SGLang/vLLM/remote have no K3 implementation), context and tensor parallelism (K3 owns CP end to end and declares supports_tp=False), pipeline parallelism (aux capture hooks one non-pipelined forward), sequence packing (K3 owns its packed-attention path), and P-EAGLE.

The target's expert dispatcher defaults to torch, the dispatcher K3's own SFT reference config runs expert parallelism with; deepep/hybridep stay opt-in through target_dispatcher since they need a matching DeepEP build.

Also here

• freeze_embeddings() on the shared MLA draft base. The recipe calls it for every draft, and it was missing, so freeze_embeddings: true raised AttributeError on the DeepSeek draft too. It matters for K3: the draft copies a 163840 x 7168 embedding table.
• The packing error message now names the actual draft class instead of always saying DeepseekV3Eagle3DraftModel.

Testing

• tests/unit_tests/speculative/test_eagle3_kimi_k3_draft.py (16 tests): forward/TTT/packing shapes, NoPE position invariance, the output gate, backbone composition, the GQA and mla_use_nope rejections, the vocab-remap and embedding-copy API, registry resolution for both architectures, and end-to-end training through the shared Eagle3TrainerModule.
• tests/unit_tests/recipes/llm/test_eagle3_kimi_k3_target.py (16 tests): backend defaults and overrides, every gate, and the resulting from_config call including the untouched draft base config.
• tests/unit_tests/speculative/ and tests/unit_tests/recipes/llm/ pass locally: the failure list is identical to the one main produces on its own (25 test_benchmark.py failures from a local environment gap, unrelated to this PR).
• No GPU run on my side: I have no Kimi K3 target weights, so everything above is CPU-verified only. As on feat(speculative): add Kimi K3 DSpark training #3263, asking @HuiyingLi to run the GPU validation; see the comment below for what a GPU run would cover that CPU cannot.

Re-targeting onto main

main now carries the merged K3 model, so this branch merges main and keeps main's version of every K3 file. One fix was needed on top:

• The non-colocated backends call _setup_online_target with draft_base_config=None, so the K3 gate reading .model_type unconditionally raised AttributeError before the backend dispatch. It now reads through getattr. This is covered by main's test_eagle3_vllm_backend.py / test_eagle3_sglang_backend.py, which fail without it.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/claude review

[HuiyingLi]

/ok to test b1f4532

[HuiyingLi]

/ok to test ac45c67