Skip to content

feat(models): add Kimi K2 + K2.7 (text-only) support#3065

Draft
faresobeid wants to merge 1 commit into
mainfrom
feat/kimi-k2-support
Draft

feat(models): add Kimi K2 + K2.7 (text-only) support#3065
faresobeid wants to merge 1 commit into
mainfrom
feat/kimi-k2-support

Conversation

@faresobeid

Copy link
Copy Markdown
Contributor

Summary

  • Kimi K2 reuses DeepSeek-V3's architecture verbatim (MLA + MoE hyperparameters, weight-key naming — Moonshot's own checkpoints ship transformers' DeepseekV3ForCausalLM as their auto_map target). KimiK2Config subclasses DeepseekV3Config directly rather than redefining its ~20 MLA/MoE fields.
  • New layers/mla.py: plain (dense, non-absorbed) Multi-head Latent Attention, shared infrastructure any MLA family can build on.
  • MoE HF↔prime conversion is reused as-is from glm4_moe — the per-expert gate/up/down + shared-experts + e_score_correction_bias layout matches exactly, no adaptation needed.
  • K2.7 (and K2.5/K2.6) ship as a multimodal wrapper around this same text backbone (language_model.-prefixed weights + a vision tower). Loading one strips the prefix and drops the vision tower entirely — text-only, no VLM/multimodal support.
  • Builds on a resurrected/adapted version of @samsja's earlier sami/kimi-k2 draft (add kimi k2  #1761, closed unmerged) — same core design (subclass DeepseekV3Config, reuse GLM's MoE conversion), updated for current main (renamed kimi_k25kimi_k2, fa4flash_attention_4, get_cu_seqlens_from_position_ids, MoEArgs.fp8), plus the K2.7 text-extraction addition and a dedup pass on the two MLA attention variants.

Verification

  • scripts/mini_moe.py --arch kimi_k2: HF↔prime roundtrip logits match to 1e-6, using transformers' own DeepseekV3ForCausalLM as the reference (real Kimi checkpoints use this exact class, so this validates the modeling code faithfully reimplements the same math).
  • Full trainer unit test suite: no new failures from this change (confirmed the same pre-existing failures reproduce identically on a clean, unmodified checkout on this GPU node — unrelated to this PR).
  • No real Kimi checkpoint was available to validate end-to-end against.
  • No inference-side changes: vLLM already serves real Kimi checkpoints natively, since their config.json declares architectures: ["DeepseekV3ForCausalLM"] and routes straight into vLLM's existing DeepSeek-V3 support.

🤖 Generated with Claude Code

Kimi K2 reuses DeepSeek-V3's architecture verbatim (MLA + MoE hyperparameters,
weight-key naming, even auto_map points at transformers' own DeepseekV3ForCausalLM),
so this subclasses DeepseekV3Config directly and adds a plain (non-sparse) MLA
attention layer (layers/mla.py) rather than redefining ~20 config fields. Reuses
GLM-4 MoE's existing HF<->prime conversion functions as-is since the MoE layout
(per-expert gate/up/down, shared-experts, e_score_correction_bias) matches exactly.

K2.7 (and K2.5/K2.6) ship as a multimodal wrapper with the same text backbone,
language_model.-prefixed. Loading one strips that prefix and drops the vision
tower entirely — text-only, no VLM support.

Verified via scripts/mini_moe.py's HF<->prime roundtrip check (max logits diff
1e-6, using transformers' own DeepseekV3ForCausalLM as the reference — Kimi's own
checkpoints use this exact class). No real Kimi checkpoint was available to
validate against directly. vLLM already serves real Kimi checkpoints natively
(architectures=DeepseekV3ForCausalLM routes into its existing DeepSeek-V3 support),
so no inference-side changes are needed.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants