SIGN IN SIGN UP

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