[yijingl-nvidia]

Description

This PR moves the Kimi K3 weight-loading code from modeling_kimi_linear.py into the new file tensorrt_llm/_torch/models/modeling_kimi_k3_loader.py.

No functional change.

modeling_kimi_linear.py had over 3,300 lines, and about 1,100 of them were weight loading. Now the model file keeps only the model and its forward pass. This improves code readability and maintenance. #20782 adds the Kimi K3 Model Streamer load to the new file.

This PR does not depend on the other PRs in the series.

Test Coverage

0 tests added, 3 tests changed. The tests only import or patch the moved code from the new file:

• tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py
• tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py
• tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

PRs in this series

These PRs refactor checkpoint loading and add Model Streamer loading. They merge in this order, top to bottom.

PR	What it does
#20498	Change checkpoint_io_policy from plain strings to the string enum CheckpointIOPolicy.
#20703	Add the string enum CheckpointFormat for the built-in checkpoint formats, make the checkpoint I/O helpers take only LoadFormat, and make the checkpoint I/O policy resolution easier to read.
#20704 (merged)	Put the PyTorch process-group setup in one module, distributed/initialization.py.
#20711	Make ModelLoader.load() shorter: move its MX, GMS, memory and model-hook helpers to new files in models/checkpoints/.
#20713	Remove the MX-only hooks from BaseCheckpointLoader.
#20717 (this PR)	Move Kimi K3 weight loading from modeling_kimi_linear.py to modeling_kimi_k3_loader.py.
#20781	Add checkpoint_io_policy="model_streamer", which loads a checkpoint with Run:ai Model Streamer.
#20782	Add the Model Streamer load for Kimi K3.

🤖 Generated with Claude Code

Dev Engineer Review

• Adds KimiK3WeightLoaderMixin in modeling_kimi_k3_loader.py and moves Kimi K3 checkpoint-loading code out of modeling_kimi_linear.py.
• KimiLinearForCausalLM now inherits the mixin. The stated intent is to keep model and forward-pass code in modeling_kimi_linear.py; the PR description reports no functional change.
• The extracted loader handles checkpoint mapping, sharding, FP8 loading, expert loading, and post-load finalization. Verify these paths after the move.
• Supplied CI results are mixed: one CI run failed; another CI run succeeded while its associated pipeline failed. Results for the changed tests are not established.

QA Engineer Review

• Three test files changed: test_kimi_linear_checkpoint.py, test_kimi_k3_situ_moe.py, and test_kimi_linear_helix_mappings.py. The changes update imports or the patched module path for moved loader helpers. The supplied change summary reports no test-logic changes.
• The checkpoint test file is listed in test-db/l0_b200.yml. The MoE test file is listed in test-db/l0_b200.yml and has selected tests listed in test-db/l0_b300.yml, including expert checkpoint selection, FP8 block-scaled dequantization, safetensors materialization, and streaming expert-load checks. No test-list entry was found for the Helix mapping test file. No manual-QA list entry was found.
• Coverage verdict: needs follow-up. The supplied CI results include a failed pipeline, and successful execution of the affected tests is not established.

Per-File QA Perspective

• tensorrt_llm/_torch/models/modeling_kimi_k3_loader.py: Adds the extracted loader mixin. Verify checkpoint mapping, sharding, FP8 handling, expert loading, finalization, and error paths.
• tensorrt_llm/_torch/models/modeling_kimi_linear.py: Changes KimiLinearForCausalLM to inherit the loader mixin. Verify model construction and loader method resolution.
• tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py: Updates imports for moved checkpoint helpers. The file is listed in tests/integration/test_lists/test-db/l0_b200.yml; its tests cover checkpoint planning, sharding, and KDA/MLA loading behavior.
• tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py: Updates imports for moved loader helpers. The file is listed in test-db/l0_b200.yml; selected expert-loading and helper tests are also listed in test-db/l0_b300.yml.
• tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py: Patches run_concurrently in the loader module for the mapping swap/restore test. No test-list entry was found.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: c0c1044
• Latest commit, not reviewed: d7c81a0

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review.

Walkthrough

Kimi K3 checkpoint-loading logic now resides in a dedicated mixin module. The causal-LM wrapper inherits the mixin, and tests import the moved loading helpers and patch the new module’s run_concurrently.

Changes

Kimi K3 Weight Loader Extraction

Layer / File(s)	Summary
Sharding and FP8 loading helpers tensorrt_llm/_torch/models/modeling_kimi_k3_loader.py, tensorrt_llm/_torch/models/modeling_kimi_linear.py	The loader module adds checkpoint sharding helpers, FP8 conversion and attention-loading routines, and parameter discovery. The model module removes the corresponding helper implementations.
Checkpoint planning and weight loading tensorrt_llm/_torch/models/modeling_kimi_k3_loader.py, tensorrt_llm/_torch/models/modeling_kimi_linear.py	The loader plans and validates checkpoint keys, loads and shards trunk weights, and loads experts from safetensors shards or the shared mapping. It checks that backend expert slots are filled.
Finalization and wrapper integration tensorrt_llm/_torch/models/modeling_kimi_k3_loader.py, tensorrt_llm/_torch/models/modeling_kimi_linear.py, tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py, tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py, tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py	The loader finalizes KDA and FP8 attention weights and applies the optional MoE MLP conversion. KimiLinearForCausalLM inherits the mixin. Tests import moved helpers and patch run_concurrently in the loader module.

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Refactor

Merge Risk: Minimal · previous review c0c10

This refactor moves Kimi K3 checkpoint-loading code into its own module and updates the tests that reference it. No functional change was identified, and no leftover references to the moved helpers were found in the checked-out code. The mergeability risk is minimal.

Pre-merge checks: 4

Check name	Status	Explanation
Title check		The title clearly and concisely identifies the main refactor: moving Kimi K3 weight loading into a dedicated loader file.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description check		The description states the purpose clearly. It moves Kimi K3 weight-loading code from modeling_kimi_linear.py to the new modeling_kimi_k3_loader.py with no functional change. It explains the reason: t…

Full details: Description check

Explanation

The description states the purpose clearly. It moves Kimi K3 weight-loading code from modeling_kimi_linear.py to the new modeling_kimi_k3_loader.py with no functional change. It explains the reason: the model file drops from over 3,300 lines and keeps only the model and forward pass. It lists the three changed tests and the series ordering table. The template's Description and Test Coverage sections are present and filled in. The 'Test Coverage' section names the changed tests but states that no new tests were added, which is acceptable for a pure move. The PR Checklist is ticked, but the description does not confirm that the author ran the existing tests, and the PR objectives report failing CI runs on recent commits. The template's API-label and title requirements are not addressed in the description. The title format was not provided for review in this context.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/TensorRT-LLM
• Branch: main
• File: .coderabbit.yaml
• Commit: f65df75
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 9db1df6c-5968-4587-9178-920094e70e95

Commits

Reviewing files that changed from the base of the PR and between 1406184 and c0c1044.

Files selected for processing: 4

• tensorrt_llm/_torch/models/modeling_kimi_k3_loader.py
• tensorrt_llm/_torch/models/modeling_kimi_linear.py
• tests/unittest/_torch/modeling/test_kimi_linear_helix_mappings.py
• tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

[mikeiovine]

Stamp on behalf of runtime devs, delegating proper review to @NVIDIA/trt-llm-models-devs

[yijingl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #79725 [ run ] triggered by Bot. Commit: 0843081 Link to invocation

[nv-xtf]

LGTM from disagg side.

[tensorrt-cicd]

PR_Github #79725 [ run ] completed with state FAILURE. Commit: 0843081
/LLM/main/L0_MergeRequest_PR pipeline #65637 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yijingl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #79965 [ run ] triggered by Bot. Commit: 1406184 Link to invocation

[tensorrt-cicd]

PR_Github #79965 [ run ] completed with state SUCCESS. Commit: 1406184
/LLM/main/L0_MergeRequest_PR pipeline #65818 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[Mgluhovskoi]

LGTM from telemetry POV.

[chienchunhung]

LGTM: This PR reorganizes Kimi K3 weight loading while preserving its behavior. The shared ModelLoader still calls model.load_weights(). Kimi now inherits that implementation from KimiK3WeightLoaderMixin in the new file. Kimi retains its specialized lazy checkpoint access, tensor sharding, FP8 handling, and expert loading.

PS: The extraction separates checkpoint loading from model execution and makes the later Model Streamer integration in #20782 easier to review.

[yijingl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #81259 [ run ] triggered by Bot. Commit: c0c1044 Link to invocation

[tensorrt-cicd]

PR_Github #81259 [ run ] completed with state FAILURE. Commit: c0c1044
/LLM/main/L0_MergeRequest_PR pipeline #66949 completed with status: 'FAILURE'

CI Report

⚠️ Resource Maintenance — Test Coverage Reduced:

Resource maintenance in LLM/main/L0_Test-x86_64-Single-GPU #9114: L40S-PyTorch-2 is temporarily paused due to Kubernetes agent provisioning and node issues.
Affected stages:
L40S-PyTorch-2

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[yijingl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #81753 [ run ] triggered by Bot. Commit: c0c1044 Link to invocation

[tensorrt-cicd]

PR_Github #81753 [ run ] completed with state SUCCESS. Commit: c0c1044
/LLM/main/L0_MergeRequest_PR pipeline #67410 completed with status: 'SUCCESS'

CI Report

Link to invocation

[yijingl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #82053 [ run ] triggered by Bot. Commit: f2bc560 Link to invocation

[tensorrt-cicd]

PR_Github #82053 [ run ] completed with state SUCCESS. Commit: f2bc560
/LLM/main/L0_MergeRequest_PR pipeline #67689 completed with status: 'SUCCESS'

CI Report

Link to invocation

[sunnyqgg]

Reviewed the move of the Kimi K3 weight loading into modeling_kimi_k3_loader.py. The moved functions and methods match the base exactly, apart from the two function-local imports of _Fp8BlockScaleWeightReadLinear. The MRO is unchanged (KimiK3WeightLoaderMixin ahead of SpecDecOneEngineForCausalLM). The external callers (runner.py, modeling_kimi_k25.py, the modeling_v2 targets) and the test monkeypatch targets still resolve, and ruff passes. Counts: 0 Blockers, 0 Possibly blocking, 2 Required, both inline: coordinate with the open PRs that edit the moved code, and correct the module docstring and the PR description about what stays in the model file.

[sunnyqgg]

[Required] Four open PRs edit the Kimi K3 load code that this PR moves. All of them are based on the same blob of modeling_kimi_linear.py (ea45af7c4f) that this PR deletes from:

• [None][perf] Read Kimi-Linear expert shards with pread into pinned host buffers #21588 replaces safe_open(...) with PinnedSafetensorsFile(...) in load_expert_file and load_split_file_expert (this line and L1073), and swaps the safe_open import in the model file.
• [None][feat] Dense DWDP for Kimi-K3 context servers and a balanced DWDP expert partition #21012 adds a dense-DWDP guard at the top of load_weights (L621).
• [TRTLLM-14978][perf] Kimi-K3 decode attention-residual kernel: CuTe DSL PDL mix with the MoE residual folded in, MXFP8 hand-off to KDA, gated/sigmoid low-M GEMM variants #21399 and [TRTLLM-14978][perf] Kimi-K3 MoE decode front: routing + moe_sort in one launch, CuTe router logits, router-first order, NVFP4 latent_down epilogue #21400 add attention.accept_prequantized_qkvg_input(fused.takes_prequantized_input) in _finalize_weight_load (around L1155-1157).
• Draft [None][chore] Add a code-filtered mypy check for tensorrt_llm and fix the errors it found #20946 adds # type: ignore[call-arg] on load_expert(*job, ...) (L1093).

Git will report conflicts in these regions. But resolving them by keeping this PR's copy of the loader as is (or, on their side, by dropping the hunk in the deleted region) silently loses the other PR's change. For #21588, nothing would catch that: its only test covers PinnedSafetensorsFile itself, not the Kimi load path.

Please link these PRs in the description and agree on the merge order. Whichever lands later has to port those hunks into modeling_kimi_k3_loader.py. When you rebase this PR over any of them, re-run the removed-vs-added comparison (e.g. git diff --color-moved) to make sure their edits end up in the loader.

[sunnyqgg]

[Required] This sentence doesn't match the code, and neither does the PR description's "the model file keeps only the model and its forward pass". The routed-expert checkpoint layouts stay in modeling_kimi_linear.py:967-1072: _load_packed_mxfp4_expert, _load_nvfp4_expert, _K3ExpertCkptSpec, _K3_EXPERT_CKPT_SPECS and _k3_expert_ckpt_spec. They read the per-expert checkpoint keys ({base}.{e}.w*.weight_packed, ...) and call quant_method.load_*_expert, and their fields are read only in this file (L595, L914-943, L960, L1006-1024). The section comment there still says "K3 streams routed experts one at a time (see load_weights)" (modeling_kimi_linear.py:970, and similarly L1136-1137), but load_weights is no longer in that file. So someone who adds a routed-expert layout, or the Model Streamer path from #20782, will look in the wrong file.

Please do one of the following:

• Move that section here. The model file already imports this module, so there is no cycle. Keep _k3_expert_ckpt_spec reachable from modeling_kimi_linear for the test that monkeypatches it, and update the tests that import the other names.
• Or say in this docstring that the per-expert layouts stay with KimiK3MoERuntime and are reached through moe.expert_ckpt_spec.

In both cases, update the load_weights pointers at modeling_kimi_linear.py:970 and :1136, and the PR description.

[sunnyqgg]

Follow-up to my review of 2026-11-04: All 9 lower-priority points I left out of it still apply at d7c81a0. They are marked [Nice to have] and none blocks the merge on its own.

• [Nice to have] No CI has run on the current head d7c81a0. The last pipeline, PR_Github #82053 (SUCCESS), ran on f2bc560, which has the same five file contents but a base 59 commits older. GitHub also reports a merge conflict with main now, so please rebase and run /bot run on the new head before merging.

[sunnyqgg]

[Nice to have] The module docstring (modeling_kimi_k3_loader.py:19-22) presents _finalize_weight_load() as just the FP8 weight-read swap. It also runs finalize_decode_weights() and _build_mtp_conv_weights() on every KDA layer (:1147-1150) and repoints q/k/v/g to the fused qkvg_proj rows (:1156-1167). A load path that skips it, even for a BF16 checkpoint, leaves _packed_conv_weight as None, and the first KDA prefill asserts (kimi_kda_mixer.py:767). Please document that, and note in the load_weights() bullet that file-grouped expert streaming needs checkpoint_dir (and checkpoint_prefix for re-prefixed keys) on the weights mapping (:990-997); without checkpoint_dir it can drop to the shared-dict fallback (:1099-1110).

[sunnyqgg]

[Nice to have] _Fp8BlockScaleWeightReadLinear stays in modeling_kimi_linear.py:723, but its only production users are now in this file (modeling_kimi_k3_loader.py:176, :1157). Because modeling_kimi_linear.py:121 imports this module, both uses need function-level imports (:171, :1142), and hoisting either one later fails with a partially-initialized-module ImportError. Consider moving the class and its helpers (modeling_kimi_linear.py:699-941) into a module under _torch/modules/, importing it at the top here, and repointing the tests that import or monkeypatch them (e.g. test_linear_output_slice.py:240). A follow-up is fine; meanwhile the comments at :169 and :1140 should say "import the class", not "the module".

[sunnyqgg]

[Nice to have] The mixin reads five host attributes, model, model_config, lm_head (modeling_kimi_k3_loader.py:785), _repurposed_tp_mapping (:713-714) and _fp8_weight_read_moe_mlp (:1171), but declares none, and the docstring names only two. The modeling_v2 K3 targets already set the private two by hand (kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/modeling.py:3246-3248, kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/modeling.py:3316-3318); a host that misses one hits an AttributeError mid-load. Please declare them as annotation-only class attributes, as MultimodalModelMixin does (modeling_multimodal_mixin.py:431), and note that the mixin must stay ahead of SpecDecOneEngineForCausalLM in the bases (modeling_kimi_linear.py:2350-2352) to override its load_weights.

[sunnyqgg]

[Nice to have] Nit: the last section header before the class, "Swap projections to the FP8 block-scale weight-read module after loading" (modeling_kimi_k3_loader.py:150-153), now also covers the checkpoint FP8 helpers and _materialize from here to :311, which run during the load rather than after it. Please add a section header above this line, e.g. "Checkpoint FP8 block-scale pairs and lazy tensors", and consider moving _has_weights (:181), which the whole loader uses, out of the swap section.

[sunnyqgg]

[Nice to have] Nit: the first definition in the new file has no comment. Please add one line, e.g. "MLA parameters derived from kv_b_proj during the load (_load_mla_fp8_kv_b() and the BF16 KV-B path); they have no checkpoint key, so _trunk_parameters() skips them." Today this is explained only in the _trunk_parameters() docstring (modeling_kimi_k3_loader.py:529-531).

[sunnyqgg]

[Nice to have] Nit: the test patches loader_mod.run_concurrently (:277, :303) but calls prod_mod.KimiLinearForCausalLM._load_trunk_params (:278, :304). The patch only takes effect because the inherited method looks up run_concurrently in the loader's globals (modeling_kimi_k3_loader.py:866), which a reader has to trace through the class bases. Call loader_mod.KimiK3WeightLoaderMixin._load_trunk_params(...) instead, or rename prod_mod to e.g. kimi_linear_mod.

[sunnyqgg]

[Nice to have] No pre-merge test reaches this streaming path. The Kimi K3 load_weights tests build no MoE layers (num_experts unset in test_kimi_k3_checkpoint.py and test_kimi_k3_mla.py; is_moe = False at test_kimi_linear_checkpoint.py:273), so expert_jobs is empty, and no test calls _load_expert_slices. The gate/up row-concat (modeling_kimi_k3_loader.py:763-764) only sees zeros (test_kimi_k3_checkpoint.py:82, test_kimi_k3_mla.py:254). Optional: drive _load_expert_slices with a stub MoE and a tmp index plus two shards (one expert split across both), asserting each slot loads once and each layer finalizes once, and add a load_weights case with distinct gate/up values at tp=2, rank=1.

[sunnyqgg]

[Nice to have] Follow-up idea: the per-expert protocol here (prepare a layer on its first expert, finalize it when the last slot lands, then check that every slot is filled and finalized; modeling_kimi_k3_loader.py:888-963, :1112-1135) lives in closures next to the two tensor sources. #20782 leaves this method as is and re-implements the same protocol in plan_model_streamer_load() (load_expert, finish). Extracting it into a small helper that takes get_tensor/get_optional_tensor would let the file-grouped, shared-dict and Model Streamer paths share one copy, so a change to the prepare/finalize contract lands once.