SIGN IN SIGN UP

feat(speculative): add Kimi K3 EAGLE-3 training (#3286)

* feat: add Kimi K3 model support

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* ci: add Kimi K3 recipe metadata

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor: keep Kimi K3 changes model-scoped

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(kimi-k3): isolate fp32 kernel parameters

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* chore(kimi-k3): remove bundled license file

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(kimi-k3): address review feedback

Avoid import-time transformer mutation, remove the dead tokenizer fallback,
and cover K3 MoE router-weight placement against reference inference for
both loop and grouped-mm expert backends.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* test(kimi-k3): cover chat effort and config resolution

Accept the documented medium thinking effort with durable input validation,
and prove both Kimi K3 model types resolve through local configs when the
installed Transformers mapping has no builtin entry.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(kimi-k3): guard optional attention residual init

Keep checkpoint-free initialization valid when attention residual mixing is
disabled while preserving pruned pipeline-stage handling, with a CPU forward
regression test for the no-residual configuration.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* chore(kimi-k3): add NVIDIA copyright notices

Keep the upstream Moonshot attribution and license while marking NVIDIA modifications so the repository copyright check recognizes both files.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(kimi-k3): simplify decoder initialization

Remove the identical MoE and dense branches around MLP weight initialization
so the shared behavior is explicit.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(speculative): add Kimi K3 EAGLE-3 training

Kimi K3 interleaves KDA layers with NoPE MLA full-attention layers, which
neither the dense EAGLE-3 draft nor the DeepSeek MLA draft (which always rotates
the rope slice) can represent. Add a dedicated draft that mirrors one K3 MLA
layer: Q-LoRA, the compressed KV latent, the unrotated rope slice, the sigmoid
output gate, K3 RMSNorm and the SiTU MLP, with KDA, the routed experts and the
attention-residual mixer left in the target.

Everything EAGLE-3 specific is inherited from the DeepSeek MLA draft through a
new _build_inner_model hook, so the TTT recurrence, the vocab-remap buffers and
the trainer are reused unchanged. freeze_embeddings, which the recipe calls for
every draft, is added to that shared MLA base.

The frozen target is loaded through AutoModel's custom expert-parallel path (K3
has no HuggingFace implementation and its 896 routed experts do not fit on one
GPU), and the combinations K3 cannot serve (non-colocated backends, context or
tensor parallelism, sequence packing, parallel drafting) are rejected before the
target is loaded.

Signed-off-by: khazic <khazzz1c@gmail.com>

* fix(speculative): gate Kimi K3 pipeline parallelism and default its dispatcher to torch

Online EAGLE-3 supervision hooks one complete target forward, so a pipelined K3
target (whose layer ModuleDict holds only its own stage, keyed by global index)
either raises a KeyError in the aux-layer capture or silently captures nothing.
The README already said pp_size must stay 1; enforce it with the other gates,
before the 93-layer target is loaded.

The expert dispatcher defaulted to hybridep, which forces the DeepEP path even
where that build is absent (the world_size==1 fallback hides it until the real
multi-node run). Default to torch, the dispatcher K3's own SFT reference config
runs expert parallelism with, and keep deepep/hybridep opt-in.

Signed-off-by: khazic <khazzz1c@gmail.com>

* fix(speculative): guard the Kimi K3 gate against a missing draft config

The non-colocated EAGLE-3 target backends call _setup_online_target with
draft_base_config=None, so reading .model_type unconditionally raised
AttributeError before the backend dispatch.

Signed-off-by: khazic <khazzz1c@gmail.com>

* docs(speculative): document the fused tensor layouts in the K3 EAGLE-3 draft

Signed-off-by: khazic <khazzz1c@gmail.com>

* fix(eagle3): reshard Kimi K3 target after forward

Signed-off-by: Huiying Li <willwin.lee@gmail.com>

---------

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
Signed-off-by: Huiying Li <willwin.lee@gmail.com>
Co-authored-by: HuiyingLi <willwin.lee@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
K
khazzz1c committed
c852b16ff97221bd97a408e224f011a817cbc8e9
Parent: 03c7e70
Committed by GitHub <noreply@github.com> on 8/11/2026, 4:03:16 AM