diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 1f1f810286ea..6a96de4a161a 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 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.`` 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" ) 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/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b3a07d3b526b..45a0128bfa0f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/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 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..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 @@ -44,6 +44,11 @@ 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) @@ -70,13 +75,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 +96,17 @@ 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 +177,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 +199,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"])) 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"):