feat(model): add Kimi K3 language backbone support - #5130
Conversation
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
Full-model MCore greedy inference validation on Lyris:
|
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
Added the Kimi K3 model verification card in commit The card follows the strict verification gates: the four persisted conversion paths, full-model forward parity, inference clean-commit rerun, and all training workflows remain unverified. It records the completed full-model direct-import/greedy-generation result as diagnostic evidence only, and documents the KDA Validation completed:
|
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
Light review — Kimi K3 bridge + MXFP4 E8M0 scales Solid, well-tested addition. The bridge follows the model-support skill patterns (local subclasses/hooks, no shared-code edits, no size-specific provider). A few observations: Correctness / logic
Docs / examples gap (per CONTRIBUTING + adding-model-support skill)
Tests — Unit coverage is good. No GPU functional conversion test for K3 yet; deferring to the verification card is reasonable given fla/TE + multi-GPU deps. Suggested test cases
No perf tests impacted (no scripts/performance/configs changes). |
| linear_attn_config[layer_list_name] = [ | ||
| layer_number | ||
| for layer_number in linear_attn_config[layer_list_name] | ||
| if layer_number <= num_hidden_layers |
There was a problem hiding this comment.
The linear_attn_config layer lists are filtered with layer_number <= num_hidden_layers, while the sibling lists just above use exclusive bounds (layer_types[:num_hidden_layers], mlp_only_layers with layer < num_hidden_layers). This is only correct if kda_layers/full_attn_layers are 1-indexed in the HF config (which matches the layer_number in config.kimi_kda_layers 1-indexed check in kimi_k3_layers.py). If they were 0-indexed, <= would incorrectly retain a layer at index num_hidden_layers. Worth a confirming comment given the mixed indexing conventions in this one function.
There was a problem hiding this comment.
Confirmed 1-indexed, so the inclusive bound is correct — and added a comment saying so in 31f8fa102.
Checked against the pinned revision 9f62e4e9: num_hidden_layers is 93, kda_layers has 69 entries spanning 1..91, full_attn_layers has 24 entries spanning 4..93. 69 + 24 = 93 and the two partition 1..93 exactly, so there is no layer 0 and the maximum is num_hidden_layers itself. The consumer side agrees — KimiK3Attention.__init__ tests layer_number in config.kimi_kda_layers against MCore's 1-indexed layer_number.
KimiK3ModelProvider subclasses MLAModelProvider, so num_floating_point_operations costed all 93 layers as full multi-latent attention, including the quadratic core-attention term. 69 of those layers are KDA linear-attention blocks. Re-cost the layers listed in kimi_kda_layers with the KDA per-layer formula (q/k/v/gate/output projections, per-head beta, low-rank forget gate, depthwise short convolutions, and the chunked delta-rule recurrence), mirroring the existing gated_delta_net override. MTP layers follow the final decoder layer's attention type, and out-of-range layer numbers raise. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Add docs/models/kimi/kimi-k3.md covering supported variants, the heterogeneous KDA/MLA schedule, conversion, and known limitations, and register it in the Kimi index and the model README. Also document why the toy-model helper filters linear_attn_config layer lists with an inclusive bound: those lists hold 1-indexed global layer numbers, unlike the 0-indexed sibling lists in the same function. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
Thanks — addressed in 1. 2. Eight new tests in 3. Docs / examples / recipe gap — added The |
|
/claude review |
|
/ok to test 31f8fa1 |
|
Light review — Kimi K3 language-backbone support Overall a clean, well-scoped model-support PR: bridge + provider + spec + custom layers/ops/pipeline are namespaced under Test coverage gap (non-blocking)
Looks correct
Suggested test cases
|
Adding a new model no longer requires editing EXPECTED_REGISTRATIONS. The manifest still guards that every listed architecture stays registered, keeps its key kind, and resolves to the same bridge class. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
The blob/main URLs 404 in link-check because the files they point at are added by this PR and do not exist on main yet. Relative links resolve on the branch, stay correct after merge, and are out of scope for lychee (which checks http/https only). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Restores the repo's standard blob/main URLs in the Kimi K3 page (reverting b3d2872) and instead excludes the four URLs from lychee, matching the existing precedent for links that do not resolve yet. The files these URLs point at are added by this same PR, so blob/main 404s until it merges. A TODO marks the excludes for removal afterwards. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
/ok to test 05fd6df |
What does this PR do?
Adds Megatron Bridge support for the Kimi K3 language backbone, including
model/provider construction, HF import, distributed checkpoint save, Megatron
greedy inference, HF export mappings, and focused model tests. This remains a
draft while strict full-checkpoint HF export and exact round-trip parity are
unverified.
Changelog
payload handling, latent-MoE support, and SiTU/KDA operators.
KimiK3ForConditionalGeneration/kimi_k3withAutoBridgeandadd HF ↔ Megatron parameter mappings.
uint8/UE8M0 tensor conversion support.A_logpaddingon import, then restore the 32 inactive zero entries on export.
vision_tower.*andmm_projector.*tensors unchanged during language-backbone export.
text_config, truncateheterogeneous layer schedules, and download only the safetensor shards
required by the selected layers.
text_configdoes notserialize dropout fields.
input_idsout of the upstream MCoreMoELayercall. The Miles forkuses this argument for hash routing, while K3 uses sigmoid routing and the
MCore version pinned by Bridge does not accept it.
Validation
Full checkpoint
moonshotai/Kimi-K3revision9f62e4e9fffbd0a83ddd60e1c209d828994b3569.iter_0000000with 48 distributed-checkpoint shards,.metadata, and the 93-layer run configurationeb91f8be:What is 15% of 240? Answer in one short sentence.15% of 240 = 0.15 × 240 = 36.<|close|>think<|sep|><|open|>response<|sep|>15% of 240 is 36.A_logaudit:eliminated all GPU-to-CPU merge fallbacks
generator did not complete within the one-hour allocation
Focused tests
uv run --no-project --with pre-commit pre-commit run --all-files: passed.0..3of9313.93112946; 769 parameters received gradients13.92106724→0.00057347in 6 steps13.56379385, last-10 mean12.54698963(-7.4965%)import, collectives, and
sharded_state_dict().Known draft gaps
parity remain unverified.
four-layer proxy reached cosine similarity
0.9998329878and Pearsoncorrelation
0.9998273849.performance tuning, and user documentation remain pending.
multimodal inference are not included; export only preserves the published
vision/projector tensors unchanged.
Original implementation references
The initial K3 modeling and mapping implementation was adapted from the
Apache-2.0 Miles Kimi K3 backend:
dc62a0bdc62a0bBefore this PR is ready for review
Additional Information