feat(speculative): add Kimi K3 DSpark training (#3263)
* 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 DSpark training Signed-off-by: khazic <khazzz1c@gmail.com> * refactor(speculative): tighten Kimi K3 DSpark draft and its docs `_expand_key_value_groups` was unreachable: `kv_b_proj` emits `num_heads * (nope + v)` columns, and the caller views that with `num_key_value_heads`, so the module only runs when the two are equal. Replaced it with an invariant assert in `__init__`, so a config that violates it fails early with a message instead of at a shape mismatch, and the view now uses `num_heads` directly (matching GLM-5.2 MLA). `build_kimi_k3_draft_config` was missing the three `model_args` preconditions its four sibling builders all check, so a YAML without `markov_head_type` raised a bare `AttributeError` from deep inside the builder. Docs and error text: the cp_size message and three comments still enumerated four model families, `_dspark_target_build` and `precompute_dspark_dist` docstrings still said DeepSeek V4 / GLM-5.2 only, and the example YAML and README carried cluster-specific paths and an RDMA workaround that belong to one site rather than to Kimi K3. Signed-off-by: khazic <khazzz1c@gmail.com> * feat(speculative): report DSpark acceptance diagnostics in validation Validation reduced only the loss, while every batch already computes the numbers that decide whether a draft is worth serving: the per-position `accept_rate@k`, the expected accepted block length `tau`, and the confidence head's calibration against measured acceptance. Training logs them over a log window; validation dropped them, so a held-out run could not answer how often the draft's proposals would actually be accepted. `_run_eval` now accumulates those numerator/denominator sums over the split and forms each ratio once after the data-parallel reduction, matching how the training window reduces them. Averaging per-batch ratios would over-weight batches with fewer measured positions, so the test asserts the ratio-of-sums result specifically. Diagnostics with a zero denominator stay absent from the dict rather than being reported as zero, which would read as collapsed acceptance and a perfectly calibrated head. Signed-off-by: khazic <khazzz1c@gmail.com> * test(speculative): cover Kimi K3 draft MLA config branches Add coverage for the two KimiK3DSparkAttention branches the existing tests never reached: the mla_use_nope=False rejection and the q_lora_rank=None path, which builds a single q_proj instead of the Q-LoRA pair and is exercised through a full forward. Signed-off-by: khazic <khazzz1c@gmail.com> * docs(speculative): document the context/noise tensor layouts in the K3 DSpark draft Signed-off-by: khazic <khazzz1c@gmail.com> * fix(dspark): 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> Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
K
khazzz1c committed
33eb83d7adcdfca252c106fe4b9b6f3cdbcac5c4
Parent: 7e6a9fb
Committed by GitHub <noreply@github.com>
on 8/10/2026, 6:10:46 PM