[TRTLLM-15177][chore] Kimi K3 post-merge cleanup: config/import/test hygiene + L0 wiring - #17413
Conversation
|
/bot run |
|
PR_Github #64579 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughKimi K3 runtime loading now supports shared FP8 conversion and staged expert loading. The test-only reference package adds MXFP4 and native SiTU MoE paths. KDA parity tests, MoE tests, and Blackwell pre-merge coverage use the updated modules. ChangesKimi K3 runtime and reference validation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant load_weights
participant _load_trunk_params
participant _load_expert_slices
participant _finalize_weight_load
load_weights->>_load_trunk_params: load trunk parameters
load_weights->>_load_expert_slices: load expert slices by shard
_load_expert_slices-->>load_weights: validate expert slots
load_weights->>_finalize_weight_load: finalize KDA and FP8 state
sequenceDiagram
participant ReferenceMoE
participant invoke_native_situ_moe
participant TRTLLMGen
ReferenceMoE->>invoke_native_situ_moe: pass routing and packed weights
invoke_native_situ_moe->>TRTLLMGen: quantize activations and invoke fused MoE
TRTLLMGen-->>invoke_native_situ_moe: return BF16 output
invoke_native_situ_moe-->>ReferenceMoE: trim padded dimensions
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py (1)
51-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer relative imports for modules in the same package.
_moe_kernelsand_mxfp4are siblings of this module. The absolute form_torch.modules.moe.kimi_k3_ref_moe.*depends ontests/unittestbeing onsys.path. A relative import removes that dependency and keeps the package self-contained. The module identity stays the same for the monkeypatch intest_kimi_k3_situ_moe.py, because that test imports the package through the same absolute root.♻️ Proposed change to relative imports
-from _torch.modules.moe.kimi_k3_ref_moe._moe_kernels import ( +from ._moe_kernels import ( assert_native_situ_supported, invoke_native_situ_moe, make_situ_alpha_beta, pack_routed_expert_weights, ) -from _torch.modules.moe.kimi_k3_ref_moe._mxfp4 import ( +from ._mxfp4 import ( DEFAULT_GROUP_SIZE, dequantize_last_dim_mxfp4, quantize_last_dim_mxfp4, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py` around lines 51 - 61, Change the imports of _moe_kernels and _mxfp4 in the kimi_k3_moe_block module to package-relative imports, while preserving the imported symbols and module identity used by test_kimi_k3_situ_moe.py.tests/integration/test_lists/test-db/l0_b200.yml (1)
93-100: 🚀 Performance & Scalability | 🔵 TrivialConfirm the tier and the pre-merge time budget for the seven added suites.
The entries are module-level, so every test in each module runs.
test_kimi_k3_situ_moe.pypacks MXFP4 expert banks and launches fused TRTLLM-Gen kernels, and the KDA parity suites run two-round fused-versus-sequential comparisons. If these entries land in the pre-merge tier, they add to the time of every PR on B200. Confirm the intended tier and measure the added wall-clock time.
[operational]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/test_lists/test-db/l0_b200.yml` around lines 93 - 100, Confirm the intended tier for all seven Kimi K3/KDA suites in the test list, keeping them out of pre-merge if their full module-level execution exceeds the B200 budget. Measure and record their aggregate wall-clock impact, including test_kimi_k3_situ_moe.py and the two-round parity suites, then update the list or tier configuration accordingly.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py`:
- Around line 2188-2191: Annotate every staged-loading helper in
tensorrt_llm/_torch/models/modeling_kimi_linear.py:2188-2191, including
_validate_checkpoint_keys, with parameterized built-in collection types for
weights and expected_keys and explicit return types; annotate weights, params,
and the checkpoint name map at 2218-2234; annotate weights and expert_jobs at
2443-2449; and preserve the explicit return type while using typed call-site
values at 2568-2571. Avoid unnecessary Any and use precise checkpoint,
parameter-map, and expert-job types throughout.
- Around line 663-665: Replace the assertion guarding cfg.latent_moe_use_norm in
the Kimi K3 runtime configuration path with an explicit ValueError when the
setting is disabled or absent, preserving the existing error message and
ensuring the check remains active under python -O.
In `@tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py`:
- Around line 41-42: Move the KimiLinearConfig and KimiKDARuntime imports below
the dependency guard and its _HAVE_DEPS evaluation so missing CUDA bindings
allow the test module to skip during collection. Keep both imports within the
guarded section used by test_fused_vs_sequential_two_rounds.
---
Nitpick comments:
In `@tests/integration/test_lists/test-db/l0_b200.yml`:
- Around line 93-100: Confirm the intended tier for all seven Kimi K3/KDA suites
in the test list, keeping them out of pre-merge if their full module-level
execution exceeds the B200 budget. Measure and record their aggregate wall-clock
impact, including test_kimi_k3_situ_moe.py and the two-round parity suites, then
update the list or tier configuration accordingly.
In `@tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py`:
- Around line 51-61: Change the imports of _moe_kernels and _mxfp4 in the
kimi_k3_moe_block module to package-relative imports, while preserving the
imported symbols and module identity used by test_kimi_k3_situ_moe.py.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 43bd6b3b-65a2-4a0a-b39b-101927d2fd8f
📒 Files selected for processing (9)
tensorrt_llm/_torch/configs/kimi_linear.pytensorrt_llm/_torch/models/modeling_kimi_linear.pytensorrt_llm/_torch/modules/kimi_k3_moe/__init__.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.pytests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.pytests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.pytests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.pytests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py
|
Reconciliation note on the reviewer suggestions to fold
No code changes from this reconciliation; the PR content and validation are unchanged. |
|
PR_Github #64579 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64604 [ run ] triggered by Bot. Commit: |
|
PR_Github #64604 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64686 [ run ] triggered by Bot. Commit: |
… B200 NaN investigation The suite's compute cases fail the cosine check with NaN on B200 in pre-merge CI while passing on GB300; keep it out of the B200 list until the machine-specific numeric issue is understood. The suite remains wired and green in the GB300 list. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…200 note Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
27317b9 to
a8f650d
Compare
|
PR_Github #64943 [ run ] triggered by Bot. Commit: |
|
PR_Github #64943 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64971 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Approve with nits.
|
PR_Github #64971 [ run ] completed with state
|
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Verified the current head preserves the staged loader ordering and FP8 swap semantics; the latest commit only restores the established list-typed config default. The open bool-to-int and PR-description count threads are non-blocking nits with no runtime or test-coverage impact.
|
/bot run |
|
PR_Github #65007 [ run ] triggered by Bot. Commit: |
|
PR_Github #65007 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65065 [ run ] triggered by Bot. Commit: |
|
PR_Github #65065 [ run ] completed with state |
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
juney-nvidia
left a comment
juney-nvidia
left a comment
There was a problem hiding this comment.
Approve from qa-function perspective.
|
/bot reuse-pipeline |
Review nit: the conversion counter accumulated the helper's bool return via implicit promotion. Return 0/1 so the accumulator's type is obvious at all call sites. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
PR_Github #65141 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #65141 [ reuse-pipeline ] completed with state |
9 participants
Summary
Post-merge cleanup follow-ups deferred from #17269 (Kimi K3 / KimiLinear model), tracked in TRTLLM-15177. Low-risk, K3-scoped hygiene.
Included
KimiLinearConfig.keys_to_ignore_at_inferencelist to tuple (immutable class-level default).modeling_kimi_linear.py: hoistgc/json/ExitStack/safetensors.safe_opento module level (removed 3 redundant localimport gc), drop thejson->_jsonalias, move thelatent_moe_use_normassert to the top ofKimiK3MoERuntime.__init__.modeling_kimi_linear.pyrefactors (behavior-neutral, code moved verbatim):KimiLinearForCausalLM.load_weightsinto focused methods (_validate_checkpoint_keys,_load_trunk_params,_load_expert_slices,_finalize_weight_load) behind a short orchestrator;_swapclosures in the MoE-MLP and MLA converters plus the inline KDAo_projconversion) into a module-level_swap_linear_to_fp8_weight_readhelper.KimiK3SparseMoeBlock+_moe_kernels/_mxfp4helpers, used only bytest_kimi_k3_situ_moe.py) out of the runtime package totests/unittest/_torch/modules/moe/kimi_k3_ref_moe/, rewriting relative imports to absolute. The runtime package keeps the gate and shared MLP/RMSNorm piecesmodeling_kimi_linear.pyuses.KimiLinearConfiginstead ofSimpleNamespace,@torch.no_grad()decorators, module-level imports, drop the__main__block, ruff-format.l0_b200.yml; the 7th (test_kda_prefill_op.py) is deliberately excluded pending the B200 prefill-op NaN (TRTLLM-15204).Deferred (still tracked in TRTLLM-15177)
kimi_k3_mlaintomodeling_kimi_linear.py: small on this base (one ~260-line module with a single importer, already built on the generalMLAclass), butkimi_k3_mla_attention.pyhas diverged on the K3 feature branch (decode-backend selection changes not yet inmain), so deleting the module here would turn the next branch sync into modify/delete conflicts. Deferred until the branches converge.kimi_k3_moeintomodeling_kimi_linear.py: [None][perf] Shard Kimi K3 MLPs with tensor parallelism #17312 simplifies the MoE package (standardGatedMLPfor the shared experts and first dense layer) and is the right base for this migration; it targets the K3 feature branch and is not inmainyet. Deferred until [None][perf] Shard Kimi K3 MLPs with tensor parallelism #17312 reachesmain.kimi_k3_ref_moe/package here vs flat*_test_utils.pymodules there). The PRs target different branches (mainvsfeat/kimi_k3), so neither blocks the other; the branch-sync merge after both land should keep [None][perf] Shard Kimi K3 MLPs with tensor parallelism #17312's placement and drop the duplicate.communication_methodplumbing was intentionally dropped. It is not dead:test_kimi_k3_situ_moe.pypassescommunication_method="ALLGATHER", andcommunication_factory.pyuses it as the default whenTRTLLM_FORCE_COMM_METHODis unset.Validation
Validated on a Blackwell node (GB300, SM103 — representative of B200; not run on B200 hardware itself):
l0_b200.ymlpass, including the changedtest_kimi_kda_fused_verify_parity.py. The excludedtest_kda_prefill_op.pyfails deterministically on B200 (NaN in every compute-running case, TRTLLM-15204); the full set of 7 passed on GB300 out-of-CI.pytest --collect-only) of all 7 suites (6 wired + 1 excluded) is clean, confirming the module-levelmodeling_kimi_linearimport is safe at collection time; imports of the relocated reference MoE module verified.Related: TRTLLM-15177
Dev Engineer Review
KimiLinearConfig.keys_to_ignore_at_inferenceto an immutable tuple.asserttoValueError.kimi_k3_mlaandkimi_k3_moeconsolidation.communication_methodplumbing unchanged.QA Engineer Review
tests/integration/test_lists/test-db/l0_b200.yml.test_kimi_kda_fused_verify_parity.pytest_kimi_kda_verify_parity.pytest_kda_cache_soundness.pytest_kda_prefill_state_parity.pytest_attn_res_op.pytest_kimi_k3_situ_moe.pytest_kda_prefill_op.pybecause its compute cases produce NaNs on B200. The exclusion references TRTLLM-15204.test_fused_vs_sequential_two_rounds.test-db/l0_b200.yml.