From 9ec5f7828900f230c3a3682cc1e3e75f5991f751 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 03:29:20 -0500 Subject: [PATCH 01/13] [TRTLLM-15177][chore] use immutable tuple for KimiLinearConfig.keys_to_ignore_at_inference Class-level mutable list default could be mutated process-wide via append; a tuple is iteration-only and cannot be shared-mutated. Deferred review nit from PR #17269. Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/configs/kimi_linear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/configs/kimi_linear.py b/tensorrt_llm/_torch/configs/kimi_linear.py index b00fa6627cd0..6e63ba9d39fc 100644 --- a/tensorrt_llm/_torch/configs/kimi_linear.py +++ b/tensorrt_llm/_torch/configs/kimi_linear.py @@ -17,7 +17,7 @@ class KimiLinearConfig(PretrainedConfig): model_type = "kimi_linear" - keys_to_ignore_at_inference = ["past_key_values"] + keys_to_ignore_at_inference = ("past_key_values",) def __init__( self, From 258ef878041604b79af7c8ef88c57831a4dbf6c8 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 03:29:47 -0500 Subject: [PATCH 02/13] [TRTLLM-15177][chore] modeling_kimi_linear.py import + validation hygiene Deferred review nits from PR #17269: - Hoist stdlib imports (gc, json, ExitStack) and safetensors.safe_open to module level; drop the three redundant local 'import gc' and the function-local json/contextlib/safetensors imports. Remove the json -> _json alias (use json.load directly). - Raise the latent_moe_use_norm precondition assert to the top of KimiK3MoERuntime.__init__, beside the routed_expert_hidden_size assert, so config validation fails before any layer is allocated. Signed-off-by: Brian Nguyen --- .../_torch/models/modeling_kimi_linear.py | 23 +++++++------------ 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 1f1f810286ea..17f251659bf7 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/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 contextlib import ExitStack from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import torch +from safetensors import safe_open from torch import nn from ..._utils import is_sm_100f @@ -391,8 +395,6 @@ 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: @@ -462,8 +464,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: @@ -608,8 +608,6 @@ 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: @@ -667,6 +665,9 @@ def __init__( assert self.moe_hidden_size is not None, ( "Kimi K3 runtime expects the latent MoE (routed_expert_hidden_size)" ) + assert getattr(cfg, "latent_moe_use_norm", False), ( + "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 +764,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( @@ -2493,13 +2491,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: From 06dc2e16cace1b026867f3bd8ff16ee44f86594b Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 03:31:59 -0500 Subject: [PATCH 03/13] [TRTLLM-15177][chore] tidy KDA fused-verify parity test Deferred review nits from PR #17269: - Build the runtime config from the real KimiLinearConfig instead of a SimpleNamespace mock. - Hoist the KimiKDARuntime import (and the new config import) to module level. - Replace the two 'with torch.no_grad():' blocks with @torch.no_grad() decorators on _make_runtime and the test, dropping an indent level. - Drop the 'if __name__ == "__main__"' runner (pytest-only). The per-layer replay-cache mocks stay SimpleNamespace: they mirror the cache-manager's KDA slot allocation, which has no standalone class to instantiate in a unit test. Signed-off-by: Brian Nguyen --- .../test_kimi_kda_fused_verify_parity.py | 123 +++++++++--------- 1 file changed, 62 insertions(+), 61 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py index 37c6469b4ce6..64a638fe53c4 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py @@ -38,6 +38,9 @@ import pytest import torch +from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime + _HAVE_DEPS = True _DEP_ERR = None try: @@ -70,13 +73,18 @@ def _is_blackwell(): LB = -5.0 +@torch.no_grad() def _make_runtime(seed): - from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime - - cfg = SimpleNamespace( + # A real KimiLinearConfig (not a SimpleNamespace) so the runtime sees the + # same config surface it does in production. ``linear_attn_config`` carries + # the per-layer KDA params the runtime reads plus the (unused here) + # kda_layers/full_attn_layers schedule the config's own validation requires. + cfg = KimiLinearConfig( hidden_size=HIDDEN, rms_norm_eps=1e-5, linear_attn_config=dict( + kda_layers=[1], + full_attn_layers=[], num_heads=H, head_dim=K, short_conv_kernel_size=W, @@ -86,23 +94,22 @@ def _make_runtime(seed): ) rt = KimiKDARuntime(cfg, layer_idx=0).to("cuda") gen = torch.Generator(device="cuda").manual_seed(seed) - with torch.no_grad(): - for name, p in rt.named_parameters(): - if name.endswith("A_log"): - p.copy_( - torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.5 - ) - elif name.endswith("dt_bias"): - p.copy_( - torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.1 - ) - else: - p.copy_( - ( - torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) - * 0.03 - ).to(p.dtype) - ) + for name, p in rt.named_parameters(): + if name.endswith("A_log"): + p.copy_( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.5 + ) + elif name.endswith("dt_bias"): + p.copy_( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.1 + ) + else: + p.copy_( + ( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) + * 0.03 + ).to(p.dtype) + ) # The fused-verify conv constants are prebuilt at weight-load finalize # time in production; the runtime never computes them lazily. Mirror # that here (after the random init above, which they snapshot). @@ -173,6 +180,7 @@ def _rep(name, a, b): return cos > 0.999 and rel < 3e-2 +@torch.no_grad() def test_fused_vs_sequential_two_rounds(): torch.manual_seed(0) B = 4 @@ -194,46 +202,39 @@ def tokens(scale=0.5): ).to(torch.bfloat16) ok = True - with torch.no_grad(): - # ---- Round 1 (no pending drafts) ---- - x1 = tokens() - out1_seq = rt._forward_verify_sequential( - x1, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices - ) - out1_fused = rt._forward_verify( - x1, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices - ) - print("round 1:") - ok &= _rep("out", out1_fused, out1_seq) - - # ---- Acceptance: 0, 1, 2, 0 drafts across the 4 requests ---- - accept = torch.tensor([0, 1, 2, 0], dtype=torch.long, device="cuda") - _promote_sequential(cache_seq, conv_pool_seq, ssm_pool_seq, accept) - cache_fused.prev_num_accepted_tokens.copy_(accept.to(torch.int32)) - - # ---- Round 2 (fused path replays the accepted drafts) ---- - x2 = tokens() - out2_seq = rt._forward_verify_sequential( - x2, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices - ) - out2_fused = rt._forward_verify( - x2, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices - ) - print("round 2 (mixed replay):") - ok &= _rep("out", out2_fused, out2_seq) - - # Committed pool state cross-check: fused pool holds the state after - # round-2's golden token; reproduce it in the sequential world by - # promoting with accept=0 (golden only). - _promote_sequential( - cache_seq, conv_pool_seq, ssm_pool_seq, torch.zeros(B, dtype=torch.long, device="cuda") - ) - ok &= _rep("committed ssm", ssm_pool_fused, ssm_pool_seq) + # ---- Round 1 (no pending drafts) ---- + x1 = tokens() + out1_seq = rt._forward_verify_sequential( + x1, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices + ) + out1_fused = rt._forward_verify( + x1, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices + ) + print("round 1:") + ok &= _rep("out", out1_fused, out1_seq) + + # ---- Acceptance: 0, 1, 2, 0 drafts across the 4 requests ---- + accept = torch.tensor([0, 1, 2, 0], dtype=torch.long, device="cuda") + _promote_sequential(cache_seq, conv_pool_seq, ssm_pool_seq, accept) + cache_fused.prev_num_accepted_tokens.copy_(accept.to(torch.int32)) + + # ---- Round 2 (fused path replays the accepted drafts) ---- + x2 = tokens() + out2_seq = rt._forward_verify_sequential( + x2, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices + ) + out2_fused = rt._forward_verify( + x2, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices + ) + print("round 2 (mixed replay):") + ok &= _rep("out", out2_fused, out2_seq) + + # Committed pool state cross-check: fused pool holds the state after + # round-2's golden token; reproduce it in the sequential world by + # promoting with accept=0 (golden only). + _promote_sequential( + cache_seq, conv_pool_seq, ssm_pool_seq, torch.zeros(B, dtype=torch.long, device="cuda") + ) + ok &= _rep("committed ssm", ssm_pool_fused, ssm_pool_seq) assert ok - - -if __name__ == "__main__": - import sys - - sys.exit(pytest.main([__file__, "-v", "-x", "-s"])) From 22d9deea88782232983252960f04db11f8e4e651 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 03:36:35 -0500 Subject: [PATCH 04/13] [TRTLLM-15177][test] wire Kimi K3 unit tests into l0_b200 Register the K3 (KimiLinear) unit suites in the single-GPU Blackwell pre-merge list so they run in CI: KDA modeling parity tests, the kimi_kda module suites, the attn-res op test, and the SiTU MoE parity test. Deferred from PR #17269 (tests shipped but absent from any list). qa/ flat-list enablement is tracked separately under TRTLLM-15036. Signed-off-by: Brian Nguyen --- tests/integration/test_lists/test-db/l0_b200.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b3a07d3b526b..07d1f550b405 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -90,6 +90,14 @@ 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 + - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py + - 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 From 8f0086bd1380b86b09da151a700a32f27243d6f5 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 03:43:44 -0500 Subject: [PATCH 05/13] [TRTLLM-15177][chore] apply ruff-format to KDA parity test Signed-off-by: Brian Nguyen --- .../modeling/test_kimi_kda_fused_verify_parity.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py index 64a638fe53c4..514f646a12ab 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py @@ -96,19 +96,14 @@ def _make_runtime(seed): gen = torch.Generator(device="cuda").manual_seed(seed) for name, p in rt.named_parameters(): if name.endswith("A_log"): - p.copy_( - torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.5 - ) + p.copy_(torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.5) elif name.endswith("dt_bias"): - p.copy_( - torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.1 - ) + p.copy_(torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.1) else: p.copy_( - ( - torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) - * 0.03 - ).to(p.dtype) + (torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.03).to( + p.dtype + ) ) # The fused-verify conv constants are prebuilt at weight-load finalize # time in production; the runtime never computes them lazily. Mirror From b7da478d35e902bbfe75e52e32aca7d22b889575 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 02:08:53 -0700 Subject: [PATCH 06/13] [TRTLLM-15177][chore] split Kimi K3 load_weights and extract the FP8 swap helper Behavior-neutral refactors deferred from #17269: - Extract the duplicated FP8 weight-read module swap (nested _swap in the MoE-MLP and MLA converters, plus the inline KDA o_proj conversion) into a module-level _swap_linear_to_fp8_weight_read helper. - Split the ~450-line KimiLinearForCausalLM.load_weights into focused methods: _validate_checkpoint_keys, _load_trunk_params, _load_expert_slices, and _finalize_weight_load, with load_weights as a short orchestrator. Code moved verbatim; no functional changes. Signed-off-by: Brian Nguyen --- .../_torch/models/modeling_kimi_linear.py | 130 ++++++++++-------- 1 file changed, 74 insertions(+), 56 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 17f251659bf7..e0a6562dae32 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -384,6 +384,27 @@ 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,), +) -> bool: + """Replace ``parent.`` with an FP8 weight-read module if it is a + plain linear of one of ``linear_types``; return whether it converted. + + 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 False + setattr(parent, attr, _Fp8BlockScaleWeightReadLinear.from_linear(child)) + child.weight.data = child.weight.data.new_empty(0) + return True + + def _convert_moe_mlps_to_fp8_weight_read( model: nn.Module, include_fused_gate_up: bool = True ) -> int: @@ -397,18 +418,6 @@ def _convert_moe_mlps_to_fp8_weight_read( """ 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: @@ -426,9 +435,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 @@ -509,11 +518,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() @@ -610,18 +615,6 @@ def _convert_mla_projections_to_fp8_weight_read(model: nn.Module) -> int: """ 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. @@ -633,7 +626,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() @@ -2198,25 +2193,19 @@ def checkpoint_name_plan(self, prefix: str): return name_map, expected_keys, expert_jobs def load_weights(self, weights: Dict): - from .modeling_utils import run_concurrently - 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, expected_keys, 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 @@ -2243,6 +2232,23 @@ def load_weights(self, weights: Dict): f"checkpoint keys, e.g. {surprising[:10]}" ) + def _load_trunk_params(self, weights: Dict, params, 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.""" + from .modeling_utils import run_concurrently + + # 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 @@ -2443,6 +2449,21 @@ 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, expert_jobs) -> None: + """Load the rank-local MXFP4 expert slices of every MoE layer into + the backend expert slots, then verify every slot was filled.""" + from .modeling_utils import run_concurrently + + device = next(self.parameters()).device + def load_expert( moe: KimiK3MoERuntime, base: str, local_slot_id: int, expert_idx: int, get_tensor ): @@ -2472,13 +2493,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 @@ -2568,6 +2582,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 @@ -2600,8 +2618,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" ) From 559bef5ff5f87e197e4b49bbbfe8b8c083d63f9c Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 02:09:07 -0700 Subject: [PATCH 07/13] [TRTLLM-15177][chore] move the test-only Kimi K3 reference MoE block out of the runtime package KimiK3SparseMoeBlock (the HF-parity reference used only by test_kimi_k3_situ_moe.py) and its _moe_kernels/_mxfp4 helpers move from tensorrt_llm/_torch/modules/kimi_k3_moe/ to tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/, with relative imports rewritten to absolute ones. The runtime package keeps the gate and the shared MLP/RMSNorm blocks that modeling_kimi_linear.py uses. Signed-off-by: Brian Nguyen --- .../_torch/modules/kimi_k3_moe/__init__.py | 33 ++++--------------- .../moe/kimi_k3_ref_moe}/_moe_kernels.py | 0 .../modules/moe/kimi_k3_ref_moe}/_mxfp4.py | 0 .../moe/kimi_k3_ref_moe}/kimi_k3_moe_block.py | 29 +++++++++++----- .../modules/moe/test_kimi_k3_situ_moe.py | 20 +++++------ 5 files changed, 37 insertions(+), 45 deletions(-) rename {tensorrt_llm/_torch/modules/kimi_k3_moe => tests/unittest/_torch/modules/moe/kimi_k3_ref_moe}/_moe_kernels.py (100%) rename {tensorrt_llm/_torch/modules/kimi_k3_moe => tests/unittest/_torch/modules/moe/kimi_k3_ref_moe}/_mxfp4.py (100%) rename {tensorrt_llm/_torch/modules/kimi_k3_moe => tests/unittest/_torch/modules/moe/kimi_k3_ref_moe}/kimi_k3_moe_block.py (97%) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py index 151a5c97f561..7e79286eea2a 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py +++ b/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", ] diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py similarity index 100% rename from tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py rename to tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py similarity index 100% rename from tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py rename to tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py similarity index 97% rename from tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py rename to tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py index 3acc33e5b24a..a493f847fccc 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py +++ b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py @@ -1,9 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""KimiK3SparseMoeBlock — in-tree Kimi K3 sparse MoE module. +"""KimiK3SparseMoeBlock — test-only Kimi K3 sparse MoE reference. -Structural mirror of HF ``KimiSparseMoeBlock`` at -``modeling_kimi.py:806-918`` end to end: +Used by ``test_kimi_k3_situ_moe.py`` as the HF-parity reference for the +native SiTU kernel path; the serving runtime (``KimiK3MoERuntime`` in +``modeling_kimi_linear.py``) does not use it. Structural mirror of HF +``KimiSparseMoeBlock`` at ``modeling_kimi.py:806-918`` end to end: * :class:`KimiK3MoEGate` for routing (see :mod:`kimi_k3_moe_gate`). * :class:`KimiK3RoutedExpertBank` — per-expert MXFP4-packed @@ -46,17 +48,26 @@ from typing import Any, List, Optional, Tuple import torch -from torch import nn - -from ._mlp import KimiK3MLP, KimiK3RMSNorm, NonSituActivation, SituAndMul -from ._moe_kernels import ( +from _torch.modules.moe.kimi_k3_ref_moe._moe_kernels import ( assert_native_situ_supported, invoke_native_situ_moe, make_situ_alpha_beta, pack_routed_expert_weights, ) -from ._mxfp4 import DEFAULT_GROUP_SIZE, dequantize_last_dim_mxfp4, quantize_last_dim_mxfp4 -from .kimi_k3_moe_gate import KimiK3MoEGate +from _torch.modules.moe.kimi_k3_ref_moe._mxfp4 import ( + DEFAULT_GROUP_SIZE, + dequantize_last_dim_mxfp4, + quantize_last_dim_mxfp4, +) +from torch import nn + +from tensorrt_llm._torch.modules.kimi_k3_moe._mlp import ( + KimiK3MLP, + KimiK3RMSNorm, + NonSituActivation, + SituAndMul, +) +from tensorrt_llm._torch.modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate class KimiK3RoutedExpertBank(nn.Module): diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py index 9b412e1acf2e..2c19444f49ab 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py @@ -26,15 +26,15 @@ import pytest import torch -from utils.util import check_accuracy - -from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory -from tensorrt_llm._torch.modules.kimi_k3_moe import KimiK3SparseMoeBlock -from tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels import ( +from _torch.modules.moe.kimi_k3_ref_moe._moe_kernels import ( is_native_situ_supported, make_situ_alpha_beta, padded_fused_shapes, ) +from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_moe_block import KimiK3SparseMoeBlock +from utils.util import check_accuracy + +from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory from tensorrt_llm._torch.modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate from tensorrt_llm._torch.utils import ActType_TrtllmGen @@ -354,7 +354,7 @@ def test_fc1_swap_mutation_breaks_accuracy(): fused, ref = _make_block_pair(config, device) # Rebuild the fused buffers with w1/w3 swapped. - from tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels import pack_routed_expert_weights + from _torch.modules.moe.kimi_k3_ref_moe._moe_kernels import pack_routed_expert_weights swapped = pack_routed_expert_weights( w1_packed=fused.expert_bank.w3_packed, @@ -383,7 +383,7 @@ def test_swiglu_act_mutation_breaks_accuracy(): config = _K3Config() fused, ref = _make_block_pair(config, device) - import tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels as mk + import _torch.modules.moe.kimi_k3_ref_moe._moe_kernels as mk torch.manual_seed(17) x = torch.randn(1, 64, config.hidden_size, dtype=torch.bfloat16, device=device) * 0.5 @@ -394,13 +394,13 @@ def swiglu_invoke(**kwargs): kwargs["act_type"] = int(ActType_TrtllmGen.SwiGlu) return orig(**kwargs) - from tensorrt_llm._torch.modules import kimi_k3_moe + from _torch.modules.moe.kimi_k3_ref_moe import kimi_k3_moe_block - kimi_k3_moe.kimi_k3_moe_block.invoke_native_situ_moe = swiglu_invoke + kimi_k3_moe_block.invoke_native_situ_moe = swiglu_invoke try: out_fused = fused(x) finally: - kimi_k3_moe.kimi_k3_moe_block.invoke_native_situ_moe = orig + kimi_k3_moe_block.invoke_native_situ_moe = orig out_ref = ref(x) with pytest.raises(Exception, match="Mismatch percentage"): From 0e128b72deef52a01b12dee191a59f987807a19f Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 06:10:26 -0700 Subject: [PATCH 08/13] [TRTLLM-15177][chore] address review: config ValueError, loader annotations, guarded test imports - KimiK3MoERuntime: raise ValueError instead of assert for the required latent-MoE / latent_moe_use_norm configuration (stays active under -O). - Annotate the staged weight-loading helpers (weights / params / name_map / expected_keys / expert_jobs) with parameterized types and return types. - KDA fused-verify parity test: move the tensorrt_llm imports behind the optional-dependency guard so missing CUDA bindings skip instead of failing collection. Signed-off-by: Brian Nguyen --- .../_torch/models/modeling_kimi_linear.py | 39 ++++++++++++------- .../test_kimi_kda_fused_verify_parity.py | 8 ++-- 2 files changed, 31 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index e0a6562dae32..2f44c69e7978 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -78,7 +78,7 @@ import json import os from contextlib import ExitStack -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple import torch from safetensors import safe_open @@ -657,12 +657,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)" - ) - assert getattr(cfg, "latent_moe_use_norm", False), ( - "Kimi K3 runtime expects latent_moe_use_norm=True" - ) + # 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) @@ -2128,7 +2128,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 @@ -2142,7 +2142,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 @@ -2192,7 +2194,7 @@ 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): + def load_weights(self, weights: Dict[str, torch.Tensor]) -> None: prefix = "language_model." if any(k.startswith("language_model.") for k in weights) else "" params = self._trunk_parameters() name_map, expected_keys, expert_jobs = self.checkpoint_name_plan(prefix) @@ -2202,7 +2204,9 @@ def load_weights(self, weights: Dict): self._load_expert_slices(weights, expert_jobs) self._finalize_weight_load(num_params, len(expert_jobs)) - def _validate_checkpoint_keys(self, weights: Dict, expected_keys, prefix: str) -> None: + 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.""" @@ -2232,7 +2236,12 @@ def _validate_checkpoint_keys(self, weights: Dict, expected_keys, prefix: str) - f"checkpoint keys, e.g. {surprising[:10]}" ) - def _load_trunk_params(self, weights: Dict, params, name_map: Dict[str, str]) -> int: + 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.""" @@ -2457,7 +2466,11 @@ def load_param(name: str, param: torch.nn.Parameter): ) return len(param_jobs) - def _load_expert_slices(self, weights: Dict, expert_jobs) -> None: + 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.""" from .modeling_utils import run_concurrently diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py index 514f646a12ab..810eeab4b4da 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py @@ -38,15 +38,17 @@ import pytest import torch -from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime - _HAVE_DEPS = True _DEP_ERR = None try: import cuda.bindings.driver # noqa: F401 import cutlass # noqa: F401 from fla.ops.kda import fused_recurrent_kda # noqa: F401 + + # The model module transitively imports the optional deps above, so it + # must stay behind the guard too or collection fails instead of skipping. + from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig + from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime except ImportError as e: _HAVE_DEPS = False _DEP_ERR = str(e) From ae0a542550a6a9b4154254732e16c6ec4cf27df5 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 12:18:30 -0700 Subject: [PATCH 09/13] [TRTLLM-15177][test] Unwire KDA prefill-op suite from l0_b200 pending 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 --- tests/integration/test_lists/test-db/l0_b200.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 07d1f550b405..8eb615c11850 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -94,7 +94,9 @@ l0_b200: - 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 - - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.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. - 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 From a8f650d3758c7f84c635724c2e990c5d935b46f1 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Fri, 7 Aug 2026 12:25:07 -0700 Subject: [PATCH 10/13] [TRTLLM-15204][test] Cite the B200 KDA prefill NaN ticket in the l0_b200 note Signed-off-by: Brian Nguyen --- tests/integration/test_lists/test-db/l0_b200.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 8eb615c11850..45a0128bfa0f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -96,7 +96,7 @@ l0_b200: - 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. + # 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 From fcc4a8bfce2797cf41d0ab6139b6249cc9bac337 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 10 Aug 2026 01:58:22 -0700 Subject: [PATCH 11/13] Address trivial review comments Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/configs/kimi_linear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/configs/kimi_linear.py b/tensorrt_llm/_torch/configs/kimi_linear.py index 6e63ba9d39fc..b00fa6627cd0 100644 --- a/tensorrt_llm/_torch/configs/kimi_linear.py +++ b/tensorrt_llm/_torch/configs/kimi_linear.py @@ -17,7 +17,7 @@ class KimiLinearConfig(PretrainedConfig): model_type = "kimi_linear" - keys_to_ignore_at_inference = ("past_key_values",) + keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, From 8d4bebeb9108ca46ff67c6bdfa967d9f17e343f9 Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 10 Aug 2026 12:39:25 -0700 Subject: [PATCH 12/13] Address trivial review comments Signed-off-by: Brian Nguyen --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 2f44c69e7978..07e4376b42fb 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -99,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). @@ -2245,8 +2245,6 @@ def _load_trunk_params( """Load every non-expert trunk parameter concurrently (with the per-parameter TP-shard / pad / fuse conversions) and return the number of parameters loaded.""" - from .modeling_utils import run_concurrently - # 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 @@ -2473,8 +2471,6 @@ def _load_expert_slices( ) -> None: """Load the rank-local MXFP4 expert slices of every MoE layer into the backend expert slots, then verify every slot was filled.""" - from .modeling_utils import run_concurrently - device = next(self.parameters()).device def load_expert( From d97a73b811e839748f7aea0e63ebe9a2d124c6cc Mon Sep 17 00:00:00 2001 From: Brian Nguyen Date: Mon, 10 Aug 2026 16:05:17 -0700 Subject: [PATCH 13/13] Return an int from _swap_linear_to_fp8_weight_read 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 --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 07e4376b42fb..6a96de4a161a 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -388,9 +388,10 @@ def _swap_linear_to_fp8_weight_read( parent: nn.Module, attr: str, linear_types: Tuple[type, ...] = (nn.Linear,), -) -> bool: +) -> int: """Replace ``parent.`` with an FP8 weight-read module if it is a - plain linear of one of ``linear_types``; return whether it converted. + 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 @@ -399,10 +400,10 @@ def _swap_linear_to_fp8_weight_read( """ child = getattr(parent, attr, None) if not isinstance(child, linear_types): - return False + return 0 setattr(parent, attr, _Fp8BlockScaleWeightReadLinear.from_linear(child)) child.weight.data = child.weight.data.new_empty(0) - return True + return 1 def _convert_moe_mlps_to_fp8_weight_read(