SIGN IN SIGN UP

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