[TRTLLM-15177][chore] Kimi K3 post-merge cleanup: config/import/test hygiene + L0 wiring

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -74,10 +74,14 @@
 from __future__ import annotations
 
 import copy
+import gc
+import json
 import os
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from contextlib import ExitStack
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
 
 import torch
+from safetensors import safe_open
 from torch import nn
 
 from ..._utils import is_sm_100f
@@ -95,7 +99,7 @@
 from ..modules.rms_norm import RMSNorm
 from ..utils import ActType_TrtllmGen
 from .modeling_speculative import SpecDecOneEngineForCausalLM
-from .modeling_utils import DecoderModel, register_auto_model
+from .modeling_utils import DecoderModel, register_auto_model, run_concurrently
 
 # A/B escape hatch: restore nn.Linear for the K3 latent MoE projections
 # instead of the min-latency fused GEMM op (read once at import).
@@ -380,6 +384,28 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return out.reshape(out_shape)
 
 
+def _swap_linear_to_fp8_weight_read(
+    parent: nn.Module,
+    attr: str,
+    linear_types: Tuple[type, ...] = (nn.Linear,),
+) -> int:
+    """Replace ``parent.<attr>`` with an FP8 weight-read module if it is a
+    plain linear of one of ``linear_types``; return the number of modules
+    converted (0 or 1), so callers can accumulate a conversion count.
+
+    Frees the original BF16 weight storage immediately: the loader holds a
+    transient name->Parameter map that keeps it alive until load returns, so
+    without this the FP8 copy is purely additive and fragments the pool the
+    FP8 GEMM autotuner and KV-cache init need.
+    """
+    child = getattr(parent, attr, None)
+    if not isinstance(child, linear_types):
+        return 0
+    setattr(parent, attr, _Fp8BlockScaleWeightReadLinear.from_linear(child))
+    child.weight.data = child.weight.data.new_empty(0)
+    return 1
+
+
 def _convert_moe_mlps_to_fp8_weight_read(
     model: nn.Module, include_fused_gate_up: bool = True
 ) -> int:
@@ -391,22 +417,8 @@ def _convert_moe_mlps_to_fp8_weight_read(
     (MLA/KDA), the routed MXFP4 experts and the dense layer-0 MLP are left in
     BF16. Returns the number of projections converted.
     """
-    import gc
-
     count = 0
 
-    def _swap(parent: nn.Module, attr: str) -> None:
-        nonlocal count
-        child = getattr(parent, attr, None)
-        if isinstance(child, nn.Linear):
-            setattr(parent, attr, _Fp8BlockScaleWeightReadLinear.from_linear(child))
-            # Release the original BF16 weight storage now. The loader holds a
-            # transient name->Parameter map that keeps it alive until load
-            # returns, so without this the FP8 copy is purely additive and
-            # fragments the pool the FP8 GEMM autotuner and KV-cache init need.
-            child.weight.data = child.weight.data.new_empty(0)
-            count += 1
-
     for layer in model.layers:
         moe = getattr(layer, "block_sparse_moe", None)
         if moe is None:
@@ -424,9 +436,9 @@ def _swap(parent: nn.Module, attr: str) -> None:
                 else ("gate_proj", "up_proj", "down_proj")
             )
             for attr in shared_attrs:
-                _swap(shared, attr)
+                count += _swap_linear_to_fp8_weight_read(shared, attr)
         for attr in ("routed_expert_down_proj", "routed_expert_up_proj"):
-            _swap(moe, attr)
+            count += _swap_linear_to_fp8_weight_read(moe, attr)
 
     # Return the freed BF16 blocks to the driver so the raw (non-caching-
     # allocator) allocations made during executor creation succeed on the
@@ -462,8 +474,6 @@ def _convert_kda_projections_to_fp8_weight_read(model: nn.Module) -> int:
     ``o_proj`` reads the decode-kernel output (not the shared hidden) and is
     converted on its own. Returns the number of projections converted.
     """
-    import gc
-
     count = 0
 
     for layer in model.layers:
@@ -509,11 +519,7 @@ def _convert_kda_projections_to_fp8_weight_read(model: nn.Module) -> int:
 
         # o_proj reads the decode-kernel output, so it is not part of the fused
         # hidden-reading group; convert it on its own.
-        o_proj = getattr(mixer, "o_proj", None)
-        if isinstance(o_proj, nn.Linear):
-            setattr(mixer, "o_proj", _Fp8BlockScaleWeightReadLinear.from_linear(o_proj))
-            o_proj.weight.data = o_proj.weight.data.new_empty(0)
-            count += 1
+        count += _swap_linear_to_fp8_weight_read(mixer, "o_proj")
 
     if count:
         gc.collect()
@@ -608,22 +614,8 @@ def _convert_mla_projections_to_fp8_weight_read(model: nn.Module) -> int:
     ``forward``), with no FP8 dequant path. Returns the number of projections
     converted.
     """
-    import gc
-
     count = 0
 
-    def _swap(parent: nn.Module, attr: str) -> None:
-        nonlocal count
-        child = getattr(parent, attr, None)
-        if isinstance(child, (nn.Linear, TrtllmLinear)):
-            setattr(parent, attr, _Fp8BlockScaleWeightReadLinear.from_linear(child))
-            # Free the original BF16 storage now (as in the MLP/KDA conversions
-            # above): the loader's transient name->Parameter map would otherwise
-            # keep it alive until load returns, making the FP8 copy purely
-            # additive on the tight DEP16 pool.
-            child.weight.data = child.weight.data.new_empty(0)
-            count += 1
-
     for layer in model.layers:
         # MLA layers are the non-KDA layers (each layer is exactly one of the
         # two); their projections live on the KimiK3MLAAttention mixer.
@@ -635,7 +627,9 @@ def _swap(parent: nn.Module, attr: str) -> None:
         # g_proj exists only when the MLA output gate is enabled; a missing
         # attr is a safe no-op.
         for attr in ("q_a_proj", "q_b_proj", "o_proj", "g_proj"):
-            _swap(mixer, attr)
+            count += _swap_linear_to_fp8_weight_read(
+                mixer, attr, linear_types=(nn.Linear, TrtllmLinear)
+            )
 
     if count:
         gc.collect()
@@ -664,9 +658,12 @@ def __init__(
         self.num_experts = cfg.num_experts
         self.top_k = cfg.num_experts_per_token
         self.moe_hidden_size = cfg.routed_expert_hidden_size
-        assert self.moe_hidden_size is not None, (
-            "Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)"
-        )
+        # ValueError (not assert): these guard unsupported checkpoint
+        # configurations and must stay active under ``python -O``.
+        if self.moe_hidden_size is None:
+            raise ValueError("Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)")
+        if not getattr(cfg, "latent_moe_use_norm", False):
+            raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True")
 
         situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0
         situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None)
@@ -763,9 +760,6 @@ def __init__(
         self.routed_expert_up_proj = nn.Linear(
             self.moe_hidden_size, cfg.hidden_size, bias=False, dtype=dtype
         )
-        assert getattr(cfg, "latent_moe_use_norm", False), (
-            "Kimi K3 runtime expects latent_moe_use_norm=True"
-        )
         # Stock fused RMSNorm (flashinfer kernel; the no-flashinfer
         # fallback is the same fp32-variance eager math as KimiK3RMSNorm).
         self.routed_expert_norm = RMSNorm(
@@ -2135,7 +2129,7 @@ def get_model_defaults(cls, llm_args) -> dict:
     # correspondingly longer load).
     # ------------------------------------------------------------------
 
-    def _trunk_parameters(self):
+    def _trunk_parameters(self) -> Dict[str, torch.nn.Parameter]:
         """Named parameters of the trunk only. Spec-dec draft modules
         (e.g. the DFlash drafter attached by SpecDecOneEngineForCausalLM)
         live in a separate checkpoint loaded by
@@ -2149,7 +2143,9 @@ def _trunk_parameters(self):
             and not name.endswith(_KIMI_K3_MLA_DERIVED_PARAM_SUFFIXES)
         }
 
-    def checkpoint_name_plan(self, prefix: str):
+    def checkpoint_name_plan(
+        self, prefix: str
+    ) -> Tuple[Dict[str, str], Set[str], List[Tuple[int, KimiK3MoERuntime, str]]]:
         """Return ``(name_map, expected_keys, expert_jobs)``.
 
         ``name_map`` maps every model parameter name to its checkpoint key
@@ -2199,26 +2195,22 @@ def checkpoint_name_plan(self, prefix: str):
             expert_jobs.append((layer_idx, moe, base))
         return name_map, expected_keys, expert_jobs
 
-    def load_weights(self, weights: Dict):
-        from .modeling_utils import run_concurrently
-
+    def load_weights(self, weights: Dict[str, torch.Tensor]) -> None:
         prefix = "language_model." if any(k.startswith("language_model.") for k in weights) else ""
-
-        # The checkpoint stores every MLA KV-B head as interleaved [K | V]
-        # rows. Runtime keeps one DeepSeek-style [all K | all V] parameter
-        # instead, so context can project directly into the FMHA layout and
-        # absorbed decode can take zero-copy K/V views.
-        mla_mixers = [
-            layer.self_attn.mixer
-            for layer in self.model.layers
-            if not getattr(layer, "is_kda", True)
-        ]
-        mla_kv_b_mixers = {id(mixer.kv_b_proj.weight): mixer for mixer in mla_mixers}
-
         params = self._trunk_parameters()
         name_map, expected_keys, expert_jobs = self.checkpoint_name_plan(prefix)
 
-        # ---- key-set validation (both directions) ----
+        self._validate_checkpoint_keys(weights, expected_keys, prefix)
+        num_params = self._load_trunk_params(weights, params, name_map)
+        self._load_expert_slices(weights, expert_jobs)
+        self._finalize_weight_load(num_params, len(expert_jobs))
+
+    def _validate_checkpoint_keys(
+        self, weights: Dict[str, torch.Tensor], expected_keys: Set[str], prefix: str
+    ) -> None:
+        """Key-set validation (both directions): every expected key must be
+        present; unmatched checkpoint keys (beyond the expected leftovers)
+        only warn."""
         ckpt_keys = set(weights.keys())
         relevant_ckpt_keys = {
             k
@@ -2245,6 +2237,26 @@ def load_weights(self, weights: Dict):
                 f"checkpoint keys, e.g. {surprising[:10]}"
             )
 
+    def _load_trunk_params(
+        self,
+        weights: Dict[str, torch.Tensor],
+        params: Dict[str, torch.nn.Parameter],
+        name_map: Dict[str, str],
+    ) -> int:
+        """Load every non-expert trunk parameter concurrently (with the
+        per-parameter TP-shard / pad / fuse conversions) and return the
+        number of parameters loaded."""
+        # The checkpoint stores every MLA KV-B head as interleaved [K | V]
+        # rows. Runtime keeps one DeepSeek-style [all K | all V] parameter
+        # instead, so context can project directly into the FMHA layout and
+        # absorbed decode can take zero-copy K/V views.
+        mla_mixers = [
+            layer.self_attn.mixer
+            for layer in self.model.layers
+            if not getattr(layer, "is_kda", True)
+        ]
+        mla_kv_b_mixers = {id(mixer.kv_b_proj.weight): mixer for mixer in mla_mixers}
+
         device = next(self.parameters()).device
 
         # MLP TP shard index (used only when a param's checkpoint shape is a
@@ -2445,6 +2457,23 @@ def load_param(name: str, param: torch.nn.Parameter):
                 )
             param.data.copy_(src.to(param.dtype))
 
+        param_jobs = [(name, params[name]) for name in name_map]
+        run_concurrently(load_param, param_jobs, num_workers=8)
+
+        logger.info(
+            f"Kimi K3: loaded {len(mla_mixers)} MLA KV-B projections in grouped runtime layout"
+        )
+        return len(param_jobs)
+
+    def _load_expert_slices(
+        self,
+        weights: Dict[str, torch.Tensor],
+        expert_jobs: List[Tuple[int, KimiK3MoERuntime, str]],
+    ) -> None:
+        """Load the rank-local MXFP4 expert slices of every MoE layer into
+        the backend expert slots, then verify every slot was filled."""
+        device = next(self.parameters()).device
+
         def load_expert(
             moe: KimiK3MoERuntime, base: str, local_slot_id: int, expert_idx: int, get_tensor
         ):
@@ -2474,13 +2503,6 @@ def load_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str):
                     lambda key: _materialize(weights[key]),
                 )
 
-        param_jobs = [(name, params[name]) for name in name_map]
-        run_concurrently(load_param, param_jobs, num_workers=8)
-
-        logger.info(
-            f"Kimi K3: loaded {len(mla_mixers)} MLA KV-B projections in grouped runtime layout"
-        )
-
         # ---- backend expert slots: file-grouped streaming ----
         # The shared lazy ``weights`` dict keeps every shard mmapped for the
         # whole load, so pages it touches cannot be dropped until the load
@@ -2493,13 +2515,8 @@ def load_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str):
         ckpt_dir = getattr(self.model_config.pretrained_config, "_name_or_path", None)
         index_path = os.path.join(ckpt_dir or "", "model.safetensors.index.json")
         if expert_jobs and ckpt_dir and os.path.isfile(index_path):
-            import json as _json
-            from contextlib import ExitStack
-
-            from safetensors import safe_open
-
             with open(index_path) as f:
-                weight_map = _json.load(f)["weight_map"]
+                weight_map = json.load(f)["weight_map"]
             per_file: Dict[str, list] = {}
             split_file_jobs = []
             for layer_idx, moe, base in expert_jobs:
@@ -2575,6 +2592,10 @@ def get_tensor(key):
                 )
             backend._weights_transformed = False
 
+    def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None:
+        """Post-load finalization: build the KDA decode fast-path constants
+        and apply the FP8 weight-read conversions (all behind their env
+        switches)."""
         # FP8 weight-read master switch (see the conversion block below).
         # The KDA conversion replaces the decode in-projection GEMV with a
         # fused FP8 qkvg GEMM in the mixer decode path, so when it is enabled
@@ -2607,8 +2628,8 @@ def get_tensor(key):
                 # unconditionally; three small fp32 tensors per layer.
                 layer.self_attn._build_mtp_conv_weights()
         logger.info(
-            f"Kimi K3: loaded {len(param_jobs)} parameters and the expert "
-            f"slices of {len(expert_jobs)} MoE layers; fused decode "
+            f"Kimi K3: loaded {num_params} parameters and the expert "
+            f"slices of {num_moe_layers} MoE layers; fused decode "
             f"in-projections on {num_kda_fused} KDA layers"
         )
 

tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py

@@ -1,37 +1,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-"""Kimi K3 sparse MoE in-tree module.
+"""Kimi K3 sparse MoE runtime pieces.
 
-Ships the promoted ``KimiK3SparseMoeBlock`` and its supporting pieces
-(``KimiK3MoEGate``, latent projections, shared experts, MXFP4-packed
-routed expert bank, native TRTLLM-Gen SiTU dispatch). Structurally
-mirrors HF ``KimiSparseMoeBlock`` at ``modeling_kimi.py:806-918``.
-
-Two kernel paths coexist under one module class:
-
-* ``use_fused_cubin=False`` — Python fallback with MXFP4 group-32
-  routed expert weights, dequantized to canonical fp32 on demand.
-  Byte-exact HF parity under random weights.
-* ``use_fused_cubin=True`` — native in-tree
-  ``torch.ops.trtllm.mxe4m3_mxe2m1_block_scale_moe_runner`` invocation
-  (``act_type=SiTu``) on checkpoint-derived MXFP4 weights shared with
-  the fallback bank. Routing goes through the op's
-  ``topk_weights``/``topk_ids`` bypass fed by the real K3 gate.
+Ships the routing gate (``KimiK3MoEGate``) and the shared MLP /
+RMSNorm building blocks (``_mlp``) used by the serving runtime
+(``KimiK3MoERuntime`` in ``modeling_kimi_linear.py``). The test-only
+HF-parity reference block (``KimiK3SparseMoeBlock`` and its MXFP4 /
+kernel helpers) lives with its test at
+``tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/``.
 """
 
-from .kimi_k3_moe_block import (
-    KimiK3RoutedExpertBank,
-    KimiK3SparseMoeBlock,
-    MoEBlockProvenance,
-    copy_hf_moe_block_weights,
-)
 from .kimi_k3_moe_gate import KimiK3MoEGate, copy_hf_moe_gate_weights
 
 __all__ = [
     "KimiK3MoEGate",
-    "KimiK3RoutedExpertBank",
-    "KimiK3SparseMoeBlock",
-    "MoEBlockProvenance",
     "copy_hf_moe_gate_weights",
-    "copy_hf_moe_block_weights",
 ]

tests/integration/test_lists/test-db/l0_b200.yml

@@ -90,6 +90,16 @@ l0_b200:
   - unittest/llmapi/test_deepseek_v4_tokenizer.py
   - unittest/_torch/modules/test_mhc.py
   - unittest/_torch/modules/test_engram.py
+  # ------------- Kimi K3 (KimiLinear) unit tests ---------------
+  - unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py
+  - unittest/_torch/modeling/test_kimi_kda_verify_parity.py
+  - unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py
+  # test_kda_prefill_op.py is NOT wired here: its compute cases produce
+  # NaN on B200 (cosine check fails; passes on GB300). Under
+  # investigation before enabling on this machine type (TRTLLM-15204).
+  - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
+  - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py
+  - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py
   - unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py TIMEOUT (15)
   # ------------- modules (non-MoE) ---------------
   - unittest/_torch/modules/test_mla_helix.py