[TRTLLM-15177][chore] Kimi K3: split long weight loaders + test hygiene

[brnguyen2]

[TRTLLM-15177][chore] Kimi K3: split long weight loaders + test hygiene

Description

Closes out TRTLLM-15177, the cleanup items deferred from the #17269 review. These are the last live items of that list — every other item either already landed on main independently or is being dropped as will-not-do.

• Split the two long weight-loader methods in tensorrt_llm/_torch/models/modeling_kimi_linear.py (review item 2.1). _load_trunk_params (~190 lines) and _load_expert_slices (~225 lines) are split into cohesive private helpers, with the code moved verbatim (no behavior change):

  • _load_trunk_params → _build_mla_head_shard_maps, _kda_tp_shard_info, _load_gate_up_fused_param, _load_mla_kv_b_param, _slice_a_log_checkpoint_pad
  • _load_expert_slices → _stream_expert_slices_by_file, _verify_expert_slots_filled

  The nested closures that carry shared load state (load_param, ensure_prepared, maybe_finalize_layer, load_expert) stay nested; lifting them would have meant threading lock/bookkeeping state through every signature for no readability gain.

• Hoist an in-function import in tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py (item 4.4): with_multi_stream moves from inside the test to the module-level guarded import block, preserving the skip-on-missing-dependency behavior.

• Use the real cache object in the parity test (item 4.1): the fused/sequential per-layer caches are now real PythonMambaCacheManager.SpeculativeState instances instead of types.SimpleNamespace mocks, so the test exercises the class's actual field set and its has_kda_replay_caches property. The unused base conv/temporal fields are None (the test passes the pools to forward_verify directly, as production does).

Test Coverage

Pure refactor plus test hygiene — no new code paths. The moved loader code is exercised unchanged by the existing Kimi K3 L0 suites; test_kimi_kda_fused_verify_parity.py covers the fused-verify path (requires Blackwell). Verified locally with pre-commit run on both files (all hooks pass) and AST/compile checks.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why.
• PR title follows the format: [JIRA ticket/NVBugs ID/GitHub issue/None][type] Summary of the change
• Test cases are provided for new code paths (see test instructions)
• Any new dependencies have been scanned for license and vulnerabilities
• CODEOWNERS updated if ownership changes
• Documentation updated as needed
• The reviewers assigned automatically/manually are appropriate for the PR.
• Please check this after reviewing the above items as appropriate for this PR.

[brnguyen2]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #72497 [ run ] triggered by Bot. Commit: 163d570 Link to invocation

[tensorrt-cicd]

PR_Github #72497 [ run ] completed with state FAILURE. Commit: 163d570
/LLM/main/L0_MergeRequest_PR pipeline #59509 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