feat(speculative): add Kimi K3 DFlash draft training (#3287)
* 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 DFlash draft training Adds a dense Kimi K3 MLA draft to the DFlash recipe, trained against the frozen K3 text backbone. The draft keeps the target's MLA projections (Q-LoRA, compressed KV latent), output gate, fp32-variance RMSNorm, and SiTU MLP; KDA linear attention, routed and shared experts, MTP, and the attention-residual mixer stay in the target only. The draft attention subclasses the target's KimiMLAAttention so the projection layout cannot drift, and overrides only the forward: keys and values span [context | noise block], K3's MLA is NoPE so position_ids are unused, and attention runs through scaled_dot_product_attention because a DFlash block batch makes both the query and key axis sequence-scale. Per-target divergence moves into DFlashDraftSpec (build_draft_config, build_target_kwargs, attention_backends, supports_context_parallel) so the recipe carries no model names: the Qwen3 draft-config surgery moves to draft_qwen3.py, the text config is resolved with get_text_config(), and the sequence-packing gate keys on the target's own _owns_packed_attention. A pp_size>1 rejection is added for every DFlash target, since hidden-state capture needs one non-pipelined target forward. Signed-off-by: khazic <khazzz1c@gmail.com> * fix(speculative): pass a backend to the Kimi K3 DFlash draft attention Main's KimiMLAAttention now takes a BackendConfig and dispatches its forward on backend.attn. The DFlash draft subclasses it and replaces the forward, so it threads an eager backend through the decoder layers to keep the parent from building attention modules the draft never uses. Signed-off-by: khazic <khazzz1c@gmail.com> * fix(dflash): reshard Kimi K3 target after forward Signed-off-by: Huiying Li <willwin.lee@gmail.com> * fix(dflash): use modern optional annotations Signed-off-by: HuiyingLi <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
1d5e944793a7bd420dbe1a3332d9ce1e592636d8
Parent: e79b92c
Committed by GitHub <noreply@github.com>
on 8/22/2026, 11:58:40 PM