Skip to content

feat(model): add Kimi K3 language backbone support - #5130

Merged
yaoyu-33 merged 15 commits into
mainfrom
yuya/kimi-k3-bridge
Jul 30, 2026
Merged

feat(model): add Kimi K3 language backbone support#5130
yaoyu-33 merged 15 commits into
mainfrom
yuya/kimi-k3-bridge

Conversation

@yaoyu-33

@yaoyu-33 yaoyu-33 commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

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

  • Add the Kimi K3 provider, heterogeneous KDA/MLA layer spec, AttnRes pipeline
    payload handling, latent-MoE support, and SiTU/KDA operators.
  • Register KimiK3ForConditionalGeneration / kimi_k3 with AutoBridge and
    add HF ↔ Megatron parameter mappings.
  • Add MXFP4 uint8/UE8M0 tensor conversion support.
  • Validate and remove the published checkpoint's zero-only KDA A_log padding
    on import, then restore the 32 inactive zero entries on export.
  • Preserve the source checkpoint's vision_tower.* and mm_projector.*
    tensors unchanged during language-backbone export.
  • Extend the toy-HF-model helper to handle nested text_config, truncate
    heterogeneous layer schedules, and download only the safetensor shards
    required by the selected layers.
  • Keep K3 dropout disabled even though the released HF text_config does not
    serialize dropout fields.
  • Keep input_ids out of the upstream MCore MoELayer call. The Miles fork
    uses this argument for hash routing, while K3 uses sigmoid routing and the
    MCore version pinned by Bridge does not accept it.
  • Add focused unit tests and an agent-readable Kimi K3 model verification card.

Validation

Full checkpoint

  • Pinned moonshotai/Kimi-K3 revision
    9f62e4e9fffbd0a83ddd60e1c209d828994b3569.
  • GPU HF → Megatron import on 48 GB200 GPUs:
    • TP2/PP3/EP8/ETP2, BF16
    • completed in 27:01
    • persisted iter_0000000 with 48 distributed-checkpoint shards,
      .metadata, and the 93-layer run configuration
  • Megatron greedy inference on 24 GB300 GPUs at clean commit eb91f8be:
    • TP1/PP3/EP8/ETP1
    • completed 32 deterministic generation steps in 22:59
    • prompt: What is 15% of 240? Answer in one short sentence.
    • exact completion:
      15% of 240 = 0.15 × 240 = 36.<|close|>think<|sep|><|open|>response<|sep|>15% of 240 is 36.
  • Source A_log audit:
    • 69 tensors, each with 96 active entries and 32 inactive entries
    • all 2,208 inactive values are exactly zero
  • Strict full-checkpoint export is not claimed:
    • the checkpoint contains 497,220 source tensors
    • TP4/PP3/EP8/ETP4 reduced each rank to about 34–35B parameters and
      eliminated all GPU-to-CPU merge fallbacks
    • the 96-GB200 run had no traceback, OOM, or NCCL error, but the full export
      generator did not complete within the one-hour allocation
    • strict HF reload and exact tensor audit therefore remain pending

Focused tests

  • uv run --no-project --with pre-commit pre-commit run --all-files: passed.
  • Focused Kimi/quantization tests: 30 passed.
  • Export-padding and multimodal passthrough tests: 9 passed.
  • Four-layer proxy generated from the pinned checkpoint:
    • layers 0..3 of 93
    • 7 source shards, 16,402 tensors, 58,480,464,384 tensor bytes
  • TP8/EP8 import and train-step smoke on 8 H100 GPUs:
    • model construction and HF import tasks completed
    • real LM loss forward/backward and SGD update completed
    • loss 13.93112946; 769 parameters received gradients
  • TP8/EP8 RP2 learnability smokes on 8 H100 GPUs:
    • fixed-window overfit: 13.921067240.00057347 in 6 steps
    • 100-window rolling run: first-10 mean 13.56379385, last-10 mean
      12.54698963 (-7.4965%)
  • TP8/EP8/PP2 state validation on 16 H100 GPUs completed model construction,
    import, collectives, and sharded_state_dict().

Known draft gaps

  • Strict full-checkpoint Megatron → HF export, HF reload, and exact round-trip
    parity remain unverified.
  • Full HF/Megatron forward-logit correlation remains unverified; the completed
    four-layer proxy reached cosine similarity 0.9998329878 and Pearson
    correlation 0.9998273849.
  • Production pretraining/SFT/PEFT recipes, checkpoint-resume validation,
    performance tuning, and user documentation remain pending.
  • VPP is not supported yet, and KDA with CP greater than 1 remains unsupported.
  • This phase covers the language backbone. Native K3 vision/video modeling and
    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:

Before this PR is ready for review

  • Read and followed the contributor guidelines.
  • Added focused unit tests.
  • Added a model verification card with verified import and inference.
  • Complete strict full-checkpoint export and exact round-trip parity.
  • Add or update the remaining user documentation.
  • Optional-component import behavior is unchanged by this draft.

Additional Information

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 27, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

yaoyu-33 and others added 4 commits July 27, 2026 15:48
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>
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

Full-model MCore greedy inference validation on Lyris:

  • Hardware/parallelism: 8 GB300 nodes, 24 ranks, TP1 / PP3 / EP8 / ETP1.
  • Before this fix, the first forward failed at the K3 AttnRes pipeline boundary: stage-boundary payload width 7168 != (1 + 3) * 7168.
  • Enabling variable_seq_lengths makes MCore exchange the actual packed activation shape at each PP boundary (4x hidden at the first boundary and 7x hidden at the second).
  • PP2/EP2 4-layer proxy completed one greedy step successfully (Slurm 2512806, exit 0).
  • Full Kimi-K3 completed all 32 greedy argmax steps (Slurm 2512803, exit 0). Prompt: What is 15% of 240? Answer in one short sentence. Output: 15% of 240 = 0.15 × 240 = 36.<|close|>think<|sep|><|open|>response<|sep|>15% of 240 is 36.
  • Focused unit tests: 5 passed.
  • pre-commit run --all-files: passed.

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

Added the Kimi K3 model verification card in commit d4cbc9a38 at examples/model_verification_cards/kimi-k3/card.yaml.

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 A_log 96-active-head versus 128-stored-entry export gap. Training is intentionally deferred for this model-only pass.

Validation completed:

  • model-card validator passed
  • git diff --check passed
  • full pre-commit suite passed

yaoyu-33 added 4 commits July 28, 2026 09:27
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>
@yaoyu-33
yaoyu-33 marked this pull request as ready for review July 29, 2026 01:54
@claude

claude Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

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

  • quantization_utils.py: the new uint8 E8M0 branches look right — decode via ldexp(1, exp-127) and re-encode via log2().round()+127 clamped to [0,254]. Consistent bias 127. Good round-trip coverage.
  • create_hf_toy_model.py: the not-selectively_downloaded guard correctly suppresses the 'no tensors removed' error when only leading shards were downloaded. One inline note on the linear_attn_config layer-list filtering using <= while sibling lists use exclusive bounds; correct only if those lists are 1-indexed.
  • kimi_k3_bridge.py: _load_one_hf_weight requires a matching _scale whenever _packed exists and raises otherwise — good fail-loud behavior. A_log zero-padding import/export round-trip is verified in both directions.

Docs / examples gap (per CONTRIBUTING + adding-model-support skill)

  • No docs/models kimi-k3 page, no examples/models/kimi K3 entry, and no K3 recipe. The skill lists docs as required for enabling a new model. Please confirm these are intentionally deferred (model-only verification pass, as the card states) or add them.
  • flop_utils.py was not updated. KDA is a linear-attention block with different FLOPs-per-token than standard MLA; training TFLOPs/GPU will be inaccurate for K3. All training card items are unverified, so may be acceptable for now — flagging for follow-up.

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

  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_provider_bridge_configures_four_layer_proxy
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_mapping_registry_covers_kda_latent_moe_and_attn_res
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_kda_a_log_import_drops_zero_padding
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_kda_a_log_import_rejects_nonzero_padding
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_kda_a_log_export_restores_zero_padding
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_export_preserves_unconverted_multimodal_weights
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_stage_boundary_pack_unpack_and_bank_schedule
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_latent_moe_normalizes_after_combine_and_before_up_projection
  • tests/unit_tests/models/kimi/test_kimi_k3_bridge.py::test_transformer_layer_does_not_forward_input_ids_to_upstream_moe
  • tests/unit_tests/models/test_quantization_utils.py::test_quantize_dequantize_mxfp4_e2m1_packed_supports_uint8_e8m0_scales
  • tests/unit_tests/models/test_quantization_utils.py::test_mxfp4_scale_geometry_checks_logical_unpacked_shape

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@yaoyu-33 yaoyu-33 added area:model Model implementations and HF bridge logic feature New capabilities, enhancements, or enablement work full-test-suite waiting-on-customer Waiting on the original author to respond labels Jul 29, 2026
yaoyu-33 and others added 2 commits July 29, 2026 11:07
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>
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

Thanks — addressed in d708406c3 and 31f8fa102.

1. linear_attn_config layer-list bound — confirmed 1-indexed, <= is correct; added the clarifying comment. Details in the inline reply.

2. flop_utils.py / KDA — fixed rather than deferred. KimiK3ModelProvider subclasses MLAModelProvider, so num_floating_point_operations was costing all 93 layers as full MLA including the quadratic core-attention term, when 69 of them are KDA linear-attention blocks. Added a KDA branch that re-costs the layers 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 convs over q/k/v, and the chunked delta-rule recurrence over the [head_dim, head_dim] state), mirroring the existing gated_delta_net override. MTP layers inherit the final decoder layer's attention type, and out-of-range layer numbers raise.

Eight new tests in TestKimiK3KdaFlops cover the exact weighted self-attn term, MTP type inheritance, the out-of-range guard, and that an all-KDA schedule scales exactly linearly in sequence length. tests/unit_tests/training/utils/test_flop_utils.py + tests/unit_tests/models/kimi/ + tests/unit_tests/models/test_quantization_utils.py: 153 passed (nemo:26.06.01-based container).

3. Docs / examples / recipe gap — added docs/models/kimi/kimi-k3.md (supported variants, the heterogeneous KDA/MLA schedule, MXFP4 dequant, A_log padding round-trip, conversion command, and known limitations), registered in docs/models/kimi/index.md and docs/models/README.md. The page carries an explicit in-progress note pointing at the verification card.

The examples/models/kimi/k3/ scripts and the training recipe are intentionally deferred, as the PR body's Known draft gaps state: the skill's acceptance bar for conversion.sh / inference.sh is a clean round-trip on a real published model, and strict full-checkpoint export plus exact round-trip parity are still unverified here. Shipping a recipe now would also mean publishing parallelism settings that no training run has validated. Both land once export parity is closed.

@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/claude review

@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 31f8fa1

@claude

claude Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

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 models/kimi/, the FLOPs calculator is updated for the KDA hybrid schedule with matching tests, quantization changes carry a dedicated uint8 E8M0 test, and docs + verification card are honest about what is and is not verified. A few observations:

Test coverage gap (non-blocking)

  • examples/conversion/create_hf_toy_model.py gained substantial new branching: _select_hub_files (index vs single-safetensors path), _language_config (top-level vs nested text_config), the selectively_downloaded guard, and the 1-indexed linear_attn_config layer-list truncation, but has no unit test. This is the piece most likely to silently break on a future HF layout change. A small unit test over _select_hub_files and _truncate_config with a K3-shaped nested config would lock in the contract the inline comments describe.
  • kimi_k3_ops.py numerical helpers (situ_and_mul, KimiRMSNorm, attn_res_aggregate) are CPU-only and unit-testable, but currently only exercised indirectly (attn_res_aggregate is monkeypatched away in the layer test). Consider a small numerical test for at least situ_and_mul and KimiRMSNorm.

Looks correct

  • The KDA FLOPs block mirrors the existing GDN pattern (per-layer re-cost, MTP-follows-final-layer, 1-indexed validation) and the exact-formula + out-of-range tests are thorough.
  • MXFP4 uint8 E8M0 encode/decode is symmetric (ldexp/exp2 with bias 127) and the roundtrip + doubling test covers it.
  • A_log zero-padding drop-on-import / restore-on-export is guarded and tested in both directions.
  • Docs and verification card correctly scope this to the language backbone and mark export/parity/training as unverified.

Suggested test cases

  • test_kimi_k3_bridge.py::test_provider_bridge_configures_four_layer_proxy
  • test_kimi_k3_bridge.py::test_mapping_registry_covers_kda_latent_moe_and_attn_res
  • test_kimi_k3_bridge.py::test_kda_a_log_import_drops_zero_padding
  • test_kimi_k3_bridge.py::test_kda_a_log_import_rejects_nonzero_padding
  • test_kimi_k3_bridge.py::test_kda_a_log_export_restores_zero_padding
  • test_kimi_k3_bridge.py::test_export_preserves_unconverted_multimodal_weights
  • test_kimi_k3_bridge.py::test_stage_boundary_pack_unpack_and_bank_schedule
  • test_kimi_k3_bridge.py::test_latent_moe_normalizes_after_combine_and_before_up_projection
  • test_kimi_k3_bridge.py::test_transformer_layer_does_not_forward_input_ids_to_upstream_moe
  • test_quantization_utils.py::test_quantize_dequantize_mxfp4_e2m1_packed_supports_uint8_e8m0_scales
  • test_flop_utils.py::TestKimiK3KdaFlops::test_kda_layers_change_flops
  • test_flop_utils.py::TestKimiK3KdaFlops::test_kda_exact_self_attn_term
  • test_flop_utils.py::TestKimiK3KdaFlops::test_more_kda_layers_lowers_flops_at_long_context
  • test_flop_utils.py::TestKimiK3KdaFlops::test_kda_is_linear_in_sequence_length
  • test_flop_utils.py::TestKimiK3KdaFlops::test_mtp_layers_follow_final_decoder_layer_type
  • test_flop_utils.py::TestKimiK3KdaFlops::test_out_of_range_kda_layers_raise
  • Recommended new coverage: unit test for create_hf_toy_model.py _select_hub_files / _truncate_config nested-config + KDA layer-list truncation.
  • No perf tests impacted (no scripts/performance/configs/ changes).

yaoyu-33 and others added 3 commits July 29, 2026 15:27
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>
@yaoyu-33
yaoyu-33 requested a review from a team as a code owner July 29, 2026 23:45
@yaoyu-33

Copy link
Copy Markdown
Contributor Author

/ok to test 05fd6df

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:model Model implementations and HF bridge logic feature New capabilities, enhancements, or enablement work full-test-suite waiting-on-customer Waiting on the original author to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant