From 2d45284b567b20e16f53231835ca273f88c5d147 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 13 Aug 2026 20:58:01 +0800 Subject: [PATCH 1/9] [None][refactor] Refactor Kimi K3 MLP Signed-off-by: Jiagan Cheng --- .../_torch/models/modeling_kimi_linear.py | 348 +++++++++++++----- .../_torch/modules/kimi_k3_moe/__init__.py | 18 - .../_torch/modules/kimi_k3_moe/_mlp.py | 302 --------------- .../modules/kimi_k3_moe/kimi_k3_moe_gate.py | 281 -------------- tensorrt_llm/_torch/modules/situ.py | 114 ++++++ .../moe/kimi_k3_ref_moe/_moe_kernels.py | 4 +- .../modules/moe/kimi_k3_ref_moe/_mxfp4.py | 14 +- .../kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py | 71 ++++ .../moe/kimi_k3_ref_moe/kimi_k3_moe_block.py | 56 ++- .../_torch/modules/moe/test_kimi_k3_mlp.py | 230 +++++++++++- .../modules/moe/test_kimi_k3_moe_gate.py | 70 +--- .../modules/moe/test_kimi_k3_situ_and_mul.py | 3 +- .../modules/moe/test_kimi_k3_situ_moe.py | 16 +- 13 files changed, 744 insertions(+), 783 deletions(-) delete mode 100644 tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py delete mode 100644 tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py delete mode 100644 tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py create mode 100644 tensorrt_llm/_torch/modules/situ.py create mode 100644 tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 0b102e49e790..cf1065e055d9 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -50,8 +50,15 @@ routed partial sums — EP partials of whole experts, or TP partials over the intermediate shards — are all-reduced in the latent space (before ``routed_expert_norm`` / ``routed_expert_up_proj``, which are -nonlinear/linear layers applied to the full sum). ``lm_head`` uses the stock -``LMHead`` (vocab-sharded + gather), so logits are identical on all ranks. +nonlinear/linear layers applied to the full sum). When attention DP is off, +the shared experts use standard MLP TP over the model TP group: gate/up are +column-sharded and down is row-sharded. Direct MoE-TP combines the shared +hidden-width partial and routed latent partial into one all-reduce after the +two streams join, then splits them before the routed norm/up projection. +Communication-backed routed paths keep the shared ``GatedMLP`` reduction, +because their routed result is already combined. Under attention DP the +shared experts stay replicated. ``lm_head`` uses the stock ``LMHead`` +(vocab-sharded + gather), so logits are identical on all ranks. Speculative decoding: SA (suffix automaton, one-engine, draft-weight-free); the KDA/MLA runtimes implement multi-token verification with deferred @@ -76,6 +83,7 @@ import copy import gc import json +import math import os from contextlib import ExitStack from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set, Tuple @@ -89,14 +97,16 @@ from ...mapping import Mapping from ...models.modeling_utils import QuantAlgo, QuantConfig from ..attention_backend import AttentionMetadata -from ..distributed import AllReduce, AllReduceStrategy +from ..distributed import AllReduce, AllReduceParams, AllReduceStrategy +from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE from ..model_config import ModelConfig from ..modules.fused_moe import ConfigurableMoE, create_moe -from ..modules.kimi_k3_moe._mlp import KimiK3MLP, KimiK3RMSNorm -from ..modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate +from ..modules.fused_moe.routing import DeepSeekV3MoeRoutingMethod +from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear as TrtllmLinear from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..modules.situ import SituAndMul from ..utils import ActType_TrtllmGen from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -123,6 +133,10 @@ _K3_MOE_TP_ENV = "TLLM_K3_MOE_TP_SIZE" _K3_MOE_EP_ENV = "TLLM_K3_MOE_EP_SIZE" +# Route the exact Kimi RMSNorm through FlashInfer's single-kernel path. Set +# this to "0" to retain the eager reference's byte-exact rounding behavior. +_FUSED_RMSNORM = os.environ.get("KIMI_K3_FUSED_RMSNORM", "1") == "1" + if TYPE_CHECKING: from transformers import PretrainedConfig @@ -188,6 +202,108 @@ _KIMI_K3_FP8_WEIGHT_READ_GATE_UP_ENV = "KIMI_K3_FP8_WEIGHT_READ_GATE_UP" +class KimiK3MoEGate(nn.Module): + """Kimi K3 gate weights and routing method for ``ConfigurableMoE``.""" + + def __init__( + self, + config: Any, + *, + logits_gemm_dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.config = config + self.top_k = config.num_experts_per_token + self.num_experts = config.num_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.moe_router_activation_func = config.moe_router_activation_func + self.num_expert_group = getattr(config, "num_expert_group", 1) + self.topk_group = getattr(config, "topk_group", 1) + self.moe_renormalize = config.moe_renormalize + self.gating_dim = config.hidden_size + + assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( + "K3 MoE gate supports sigmoid or softmax scoring only" + ) + + # The checkpoint stores the gate weight in bf16. Storing it in bf16 + # permits the single bf16xbf16 router GEMM while retaining fp32 output. + weight_dtype = logits_gemm_dtype or torch.float32 + self.weight = nn.Parameter( + torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(self.num_experts, dtype=torch.float32, device=device) + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Compute fp32 routing logits shaped ``[num_tokens, num_experts]``.""" + hidden_2d = hidden_states.reshape(-1, self.gating_dim) + if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_2d.contiguous(), + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return torch.nn.functional.linear( + hidden_2d.type(torch.float32), + self.weight.type(torch.float32), + None, + ) + + @property + def routing_method(self) -> DeepSeekV3MoeRoutingMethod: + """Return the shared DeepSeek-V3 router used by ``ConfigurableMoE``.""" + if self.moe_router_activation_func != "sigmoid": + raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") + if not self.moe_renormalize: + raise ValueError( + "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." + ) + return DeepSeekV3MoeRoutingMethod( + top_k=self.top_k, + n_group=self.num_expert_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + is_fused=True, + ) + + +class KimiK3RMSNorm(nn.Module): + """RMSNorm matching the Kimi checkpoint implementation's rounding.""" + + def __init__( + self, + hidden_size: int, + eps: float = 1e-6, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if ( + _FUSED_RMSNORM + and IS_FLASHINFER_AVAILABLE + and hidden_states.is_cuda + and hidden_states.dtype in (torch.float16, torch.bfloat16) + and self.weight.dtype == hidden_states.dtype + ): + from ..custom_ops import flashinfer_rmsnorm + + return flashinfer_rmsnorm(hidden_states.contiguous(), self.weight, self.eps) + input_dtype = hidden_states.dtype + hidden_states_float = hidden_states.to(torch.float32) + variance = hidden_states_float.pow(2).mean(-1, keepdim=True) + hidden_states_float = hidden_states_float * torch.rsqrt(variance + self.eps) + return self.weight * hidden_states_float.to(input_dtype) + + def _resolve_fp8_weight_read_gates() -> tuple[bool, bool, bool]: """Resolve the FP8 weight-read switches into (master, kda, kda_glue). @@ -302,7 +418,7 @@ def _apply_attn_res( # --------------------------------------------------------------------------- -# Dense / shared-expert MLP: fused [gate | up] layout (``KimiK3MLP``). +# Dense / shared-expert MLP: fused [gate | up] layout (``GatedMLP``). # # The HF checkpoint stores separate ``gate_proj`` / ``up_proj`` tensors; # ``load_weights`` row-concatenates them into ``gate_up_proj`` (see @@ -328,7 +444,7 @@ def _gate_up_ckpt_keys(fused_key: str) -> Tuple[str, str]: class _Fp8BlockScaleWeightReadLinear(nn.Module): - """Bias-free ``nn.Linear`` replacement that reads its weight at FP8. + """Bias-free linear replacement that reads its weight at FP8. The BF16 weight ``[out, in]`` is quantized once (at load) to ``float8_e4m3fn`` with 128x128 block scales, then served through the @@ -353,6 +469,16 @@ def __init__( self.register_buffer("weight", weight_fp8, persistent=False) self.register_buffer("weight_scale", weight_scale, persistent=False) + @property + def has_fp8_qdq(self) -> bool: + """Match the ``Linear`` interface consumed by ``GatedMLP``.""" + return False + + @property + def has_w4a8_nvfp4_fp8(self) -> bool: + """Match the ``Linear`` interface consumed by ``GatedMLP``.""" + return False + @staticmethod def quantize_weight(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """BF16 ``[out, in]`` weight -> (FP8 weight, deep_gemm-ready scale). @@ -390,12 +516,21 @@ def quantize_weight(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: return weight_fp8, weight_scale @classmethod - def from_linear(cls, linear: nn.Linear) -> "_Fp8BlockScaleWeightReadLinear": + def from_linear(cls, linear: nn.Linear | TrtllmLinear) -> "_Fp8BlockScaleWeightReadLinear": assert linear.bias is None, "FP8 weight read expects a bias-free Linear" weight_fp8, weight_scale = cls.quantize_weight(linear.weight.data) return cls(weight_fp8, weight_scale, linear.out_features) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward( + self, + x: torch.Tensor, + *, + all_reduce_params: Optional[AllReduceParams] = None, + lora_params: Optional[dict] = None, + layer_idx: Optional[int] = None, + ) -> torch.Tensor: + if lora_params: + raise NotImplementedError("Kimi K3 FP8 weight read does not support LoRA.") out_shape = x.shape[:-1] + (self.out_features,) out = torch.ops.trtllm.fp8_swap_ab_gemm( x.reshape(-1, x.shape[-1]), @@ -448,7 +583,7 @@ def _convert_moe_mlps_to_fp8_weight_read( continue shared = getattr(moe, "shared_experts", None) if shared is not None: - # KimiK3MLP fuses gate and up into gate_up_proj; keep the split + # GatedMLP fuses gate and up into gate_up_proj; keep the split # names too so either MLP layout converts. The fused gate_up read # only pays off when attention DP re-reads it per rank per step; # under TP the bf16 GEMM overlaps on the aux stream and the FP8 @@ -459,7 +594,12 @@ def _convert_moe_mlps_to_fp8_weight_read( else ("gate_proj", "up_proj", "down_proj") ) for attr in shared_attrs: - count += _swap_linear_to_fp8_weight_read(shared, attr) + child = getattr(shared, attr, None) + if isinstance(child, TrtllmLinear) and child.tp_size != 1: + continue + count += _swap_linear_to_fp8_weight_read( + shared, attr, linear_types=(nn.Linear, TrtllmLinear) + ) for attr in ("routed_expert_down_proj", "routed_expert_up_proj"): count += _swap_linear_to_fp8_weight_read(moe, attr) @@ -714,7 +854,9 @@ def __init__( hidden_size=self.moe_hidden_size, intermediate_size=cfg.moe_intermediate_size, dtype=dtype, - reduce_results=True, + # Kimi owns the latent reduction so direct MoE-TP can combine it + # with the shared-expert partial in one collective below. + reduce_results=False, model_config=routed_moe_model_config, override_quant_config=routed_quant_config, layer_idx=layer_idx, @@ -767,28 +909,40 @@ def __init__( self.experts_per_rank = len(local_expert_ids) self.expert_lo = local_expert_ids[0] self.expert_hi = self.expert_lo + self.experts_per_rank + routed_comm = self.routed_experts.comm - # Shared experts stay replicated (DeepSeek's attention-DP - # semantics): ConfigurableMoE owns its own reduction, so there is - # no existing collective for column-shard partial sums to ride — - # the direct-path shared-expert TP (partials joining the MoE - # combine RS / routed allreduce) needs a partial-carry hook in the - # wrapper before it can be ported (follow-up). Instead their cost - # is hidden by running them on the aux stream, overlapped with the - # routed dispatch/expert/combine chain (see forward()). shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts - self.shared_experts = KimiK3MLP( + attention_dp = model_config.mapping.enable_attention_dp + shared_model_config = copy.copy(model_config) + shared_model_config.quant_config = QuantConfig() + # Under attention DP each rank owns different tokens, so the shared + # expert is replicated (TP size 1) and must not reduce across ranks. + # Direct MoE-TP leaves both branches as partials for one concatenated + # all-reduce. Communication-backed MoE already returns a complete + # routed result, so GatedMLP must reduce only its shared partial. + use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1 + self._use_combined_all_reduce = ( + use_shared_tp and routed_comm is None and self.routed_experts.all_reduce is not None + ) + self.shared_experts = GatedMLP( hidden_size=cfg.hidden_size, intermediate_size=shared_intermediate, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - use_fused_activation=True, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), dtype=dtype, + config=shared_model_config, + overridden_tp_size=1 if attention_dp else None, + reduce_output=use_shared_tp and not self._use_combined_all_reduce, + layer_idx=layer_idx, + is_shared_expert=True, ) - # Side stream (+ fork/join events) for overlapping the replicated - # shared-expert compute with the routed dispatch/expert/combine - # chain; see forward(). Only engaged when multi-stream is active - # (CUDA graphs on) and aux_stream is set, otherwise both run in + # Side stream (+ fork/join events) for overlapping shared-expert + # compute with the routed chain. Only engaged when multi-stream is + # active (CUDA graphs on) and aux_stream is set; otherwise both run in # order on the default stream. self.aux_stream = aux_stream self.moe_main_event = torch.cuda.Event() @@ -805,6 +959,15 @@ def __init__( hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype ) + def _routed_up_projection(self, hidden_states: torch.Tensor) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance( + self.routed_expert_up_proj, nn.Linear + ): + return self.routed_expert_up_proj(hidden_states) + return torch.ops.trtllm.dsv3_fused_a_gemm_op( + hidden_states, self.routed_expert_up_proj.weight.t(), None, None + ) + @staticmethod def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: """Resolve the routed-expert ``(moe_tp, moe_ep)`` split. @@ -914,6 +1077,7 @@ def forward(self, hidden_states: torch.Tensor, all_rank_num_tokens=None) -> torc """``hidden_states``: ``[num_tokens, hidden_size]`` bf16.""" identity = hidden_states router_logits = self.gate.compute_logits(hidden_states) + moe_all_reduce = self.routed_experts.all_reduce def _routed_output(): # Latent down/up projections via the min-latency fused GEMM op: @@ -939,24 +1103,21 @@ def _routed_output(): router_logits, all_rank_num_tokens=all_rank_num_tokens, ) - # EP partial latent sums are completed by the wrapper's own - # reduction BEFORE the (nonlinear) latent norm. + if self._use_combined_all_reduce: + return y + # EP partial latent sums must be completed before the nonlinear + # latent norm. Communication-backed paths combine internally; + # direct TP paths use the wrapper's AllReduce explicitly. + if self.routed_experts.comm is None and moe_all_reduce is not None: + y = moe_all_reduce(y) y = self.routed_expert_norm(y) - if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance( - self.routed_expert_up_proj, nn.Linear - ): - return self.routed_expert_up_proj(y) - return torch.ops.trtllm.dsv3_fused_a_gemm_op( - y, self.routed_expert_up_proj.weight.t(), None, None - ) + return self._routed_up_projection(y) - # Shared experts are replicated (computed once per rank) and depend - # only on the block input, not on the routed dispatch/expert/combine - # chain -- so run them on the aux stream to overlap with the serial - # EP dispatch/combine collectives. Multi-stream engages only under - # CUDA graphs; otherwise both run in order on the default stream - # with an identical result. Added after the routed combine so the - # replicated shared output is not double counted. + # Shared experts depend only on the block input, so overlap their GEMMs + # with the routed dispatch/expert/combine chain. Multi-stream engages + # only under CUDA graphs; otherwise both branches run in order on the + # default stream. Direct MoE-TP leaves both branches as partial sums + # until the streams join, then reduces them with one collective. routed_out, shared_out = maybe_execute_in_parallel( _routed_output, lambda: self.shared_experts(identity), @@ -965,6 +1126,17 @@ def _routed_output(): self.aux_stream, disable_on_compile=True, ) + if self._use_combined_all_reduce: + combined = moe_all_reduce(torch.cat((shared_out, routed_out), dim=-1)) + shared_out, routed_latent = torch.split( + combined, + (self.hidden_size, self.moe_hidden_size), + dim=-1, + ) + # The column split is a strided view; FlashInfer RMSNorm expects + # a dense last dimension. + routed_latent = self.routed_expert_norm(routed_latent.contiguous()) + routed_out = self._routed_up_projection(routed_latent) return routed_out + shared_out @@ -2074,37 +2246,35 @@ def __init__( else: situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) - # Dense-MLP TP semantics (DeepSeek _compute_mlp_tp_size - # pattern): replicated under attention-DP — each rank runs - # only its own tokens, so a weight shard would need an extra - # gather/scatter — and sharded like the shared experts - # otherwise (column gate_up, row down); the partial sums are - # all-reduced right after the call in forward(). - self._mlp_tp_size = ( - model_config.mapping.tp_size - if ( - not model_config.mapping.enable_attention_dp - and model_config.mapping.tp_size > 1 - and cfg.intermediate_size % model_config.mapping.tp_size == 0 - ) - else 1 - ) - self.mlp = KimiK3MLP( + attention_dp = model_config.mapping.enable_attention_dp + if attention_dp: + self.mlp_tp_size = 1 + else: + self.mlp_tp_size = math.gcd(cfg.intermediate_size, model_config.mapping.tp_size) + if self.mlp_tp_size > model_config.mapping.gpus_per_node: + self.mlp_tp_size = math.gcd( + self.mlp_tp_size, model_config.mapping.gpus_per_node + ) + mlp_model_config = copy.copy(model_config) + mlp_model_config.quant_config = QuantConfig() + # K3's dense layer is BF16, so a unit block size gives the same + # subgroup selection as DeepSeek-V3. Attention DP replicates the + # MLP because ranks own different tokens; otherwise the subgroup + # is block-aligned and stays within one node. + self.mlp = GatedMLP( hidden_size=cfg.hidden_size, - intermediate_size=cfg.intermediate_size // self._mlp_tp_size, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - use_fused_activation=True, + intermediate_size=cfg.intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), dtype=dtype, - ) - self._mlp_allreduce = ( - AllReduce( - mapping=model_config.mapping, - strategy=model_config.allreduce_strategy, - dtype=dtype, - ) - if self._mlp_tp_size > 1 - else None + config=mlp_model_config, + overridden_tp_size=self.mlp_tp_size, + reduce_output=self.mlp_tp_size > 1, + layer_idx=layer_idx, ) # Stock fused RMSNorm for the plain (whole-tensor) norms; numerics @@ -2177,9 +2347,6 @@ def forward( ) else: hidden_states = self.mlp(hidden_states) - if getattr(self, "_mlp_allreduce", None) is not None: - # TEP-sharded dense MLP: sum the row-parallel partials. - hidden_states = self._mlp_allreduce(hidden_states) prefix_sum = prefix_sum + hidden_states return prefix_sum, block_residual @@ -2199,8 +2366,8 @@ def __init__(self, model_config: ModelConfig): # One side stream shared across all layers. KDA overlaps its small # forget-gate projection chain with qkvg during decode and verify; - # MoE overlaps replicated shared-expert compute with routed - # dispatch/expert/combine. + # MoE overlaps shared-expert compute and its optional TP reduction + # with the routed dispatch/expert/combine chain. self.aux_stream = torch.cuda.Stream() self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, dtype=dtype) @@ -2483,11 +2650,10 @@ def _load_trunk_params( device = next(self.parameters()).device - # MLP TP shard index (used only when a param's checkpoint shape is a - # tp_size multiple of the param shape — the dense L0 MLP with - # attention-DP off; shapes match and no slicing runs otherwise). - # Every mode that shards these fused-MLP tensors shards by tp_rank. - shared_tp_rank = self.model_config.mapping.tp_rank + # MLP TP shard index. A dense MLP whose intermediate size does not + # divide model TP uses a smaller repeated TP subgroup, so its local + # shard rank is model tp_rank modulo the parameter's shard count. + model_tp_rank = self.model_config.mapping.tp_rank # KDA head-shard (attention-DP off): rank r loads head rows/cols # [r*local : (r+1)*local] of every head-major KDA tensor. kda_tp_size, kda_tp_rank = 1, 0 @@ -2517,11 +2683,11 @@ def load_param(name: str, param: torch.nn.Parameter): inter = param.shape[0] // 2 if gate.shape[0] != inter and gate.shape[0] % inter == 0: # TP-sharded fused MLP (shared experts on the direct - # MoE path, dense MLP with attention-DP off): take - # this rank's MATCHING row block from each half so - # the SiTU gate/up pairs stay aligned. shared_tp_rank - # == tp_rank in every mode that shards these. - lo = shared_tp_rank * inter + # MoE path, dense MLP with attention-DP off): take this + # subgroup rank's matching row block from each half so + # the SiTU gate/up pairs stay aligned. + shard_count = gate.shape[0] // inter + lo = (model_tp_rank % shard_count) * inter gate = gate[lo : lo + inter] up = up[lo : lo + inter] if gate.shape != (inter, param.shape[1]) or up.shape != gate.shape: @@ -2639,7 +2805,8 @@ def load_param(name: str, param: torch.nn.Parameter): and src.shape[0] % param.shape[0] == 0 and src.shape[1:] == param.shape[1:] ): - lo = shared_tp_rank * param.shape[0] + shard_count = src.shape[0] // param.shape[0] + lo = (model_tp_rank % shard_count) * param.shape[0] param.data.copy_(src[lo : lo + param.shape[0]].to(param.dtype)) return if ( @@ -2647,7 +2814,8 @@ def load_param(name: str, param: torch.nn.Parameter): and src.shape[1] % param.shape[1] == 0 and src.shape[0] == param.shape[0] ): - lo = shared_tp_rank * param.shape[1] + shard_count = src.shape[1] // param.shape[1] + lo = (model_tp_rank % shard_count) * param.shape[1] param.data.copy_(src[:, lo : lo + param.shape[1]].to(param.dtype)) return # MLA head padding (96 -> 128 query heads, see diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py deleted file mode 100644 index 7e79286eea2a..000000000000 --- a/tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 sparse MoE runtime pieces. - -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_gate import KimiK3MoEGate, copy_hf_moe_gate_weights - -__all__ = [ - "KimiK3MoEGate", - "copy_hf_moe_gate_weights", -] diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py deleted file mode 100644 index b0b0d7563c98..000000000000 --- a/tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py +++ /dev/null @@ -1,302 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Dense MLP helper for the in-tree Kimi K3 MoE block. - -Kimi K3 uses the ``situ`` activation (not SiLU/SwiGLU). ``SituAndMul`` -computes ``beta * tanh(gate / beta) * sigmoid(gate)`` on the gate half -and optionally applies ``linear_beta * tanh(up / linear_beta)`` on the -up half, then multiplies. ``KimiK3MLP`` is the fused ``gate_up_proj + -down_proj`` layout used by the shared expert stack in HF -``KimiSparseMoeBlock`` — the same shape TRT-LLM's ``GatedMLP`` uses. - -Two activation paths coexist: - -* the eager fp32 ``SituAndMul`` module — the byte-exact HF reference, - used by the parity-test MoE block and as the fallback; -* the fused Triton ``trtllm::situ_and_mul`` custom op (same fp32 math - in a single kernel, modeled on ``modules/swiglu.py``'s - ``silu_and_mul_kernel``), enabled per ``KimiK3MLP`` instance via - ``use_fused_activation=True`` (the runtime model opts in). The op is - CUDA-graph-safe: no host synchronization and no data-dependent - control flow. -""" - -from __future__ import annotations - -import os -from typing import Mapping, Optional - -import torch -import triton # type: ignore[import] -import triton.language as tl # type: ignore[import] -import triton.language.extra.libdevice as tldevice # type: ignore[import] -from torch import nn - -from ...flashinfer_utils import IS_FLASHINFER_AVAILABLE - -# Route the RMSNorm forward through flashinfer's single-kernel fused RMSNorm -# instead of the eager pow/mean/rsqrt/mul/cast chain. Set to "0" to fall back -# to the eager reference (the exact-parity rollback lever). -_FUSED_RMSNORM = os.environ.get("KIMI_K3_FUSED_RMSNORM", "1") == "1" - - -class SituAndMul(nn.Module): - """K3 SiTU activation with gate/up multiplicative gating. - - Byte-identical to HF ``modeling_kimi.py``'s ``SituAndMul`` at - lines 41-59. Runs the math in fp32 for numerical stability - (matches HF), then casts back to the input's dtype. - """ - - def __init__( - self, - *, - beta: float = 1.0, - linear_beta: Optional[float] = None, - ) -> None: - super().__init__() - self.beta = beta - self.linear_beta = linear_beta - - def forward(self, x: torch.Tensor) -> torch.Tensor: - d = x.shape[-1] // 2 - gate = x[..., :d].to(torch.float32) - up = x[..., d:].to(torch.float32) - situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) - if self.linear_beta is not None: - up = self.linear_beta * torch.tanh(up / self.linear_beta) - return (situ_a * up).to(x.dtype) - - -@triton.jit -def situ_and_mul_kernel( - o_ptr, - o_stride, - x_ptr, - x_stride, - d, - beta, - linear_beta, - BLOCK_SIZE: tl.constexpr, - HAS_LINEAR_BETA: tl.constexpr, -) -> None: - """Fused :class:`SituAndMul` on a packed ``[gate | up]`` row layout. - - Loads ``gate = x[i, :d]`` and ``up = x[i, d:2d]``, computes (fp32) - ``beta * tanh(gate / beta) * sigmoid(gate) * up'`` with - ``up' = linear_beta * tanh(up / linear_beta)`` when - ``HAS_LINEAR_BETA`` else ``up``, and stores the product rounded to - ``o_ptr``'s element type. - """ - i = tl.program_id(axis=0).to(tl.int64) - j = tl.program_id(axis=1) - - o_row_ptr = o_ptr + o_stride * i - x_row_ptr = x_ptr + x_stride * i - - offsets = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < d - - gate = tl.load(x_row_ptr + offsets, mask=mask).to(tl.float32) - up = tl.load(x_row_ptr + offsets + d, mask=mask).to(tl.float32) - - situ_a = beta * tldevice.tanh(gate / beta) * tl.sigmoid(gate) - if HAS_LINEAR_BETA: - up = linear_beta * tldevice.tanh(up / linear_beta) - result = situ_a * up - - tl.store(o_row_ptr + offsets, result, mask=mask) - - -@torch.library.custom_op("trtllm::situ_and_mul", mutates_args=()) -def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor: - """Fused SiTU activation (single Triton kernel, fp32 internal math). - - Args: - x: ``[num_tokens, 2 * d]`` packed ``[gate | up]`` GEMM output - (fp16/bf16/fp32; the last dim must be contiguous). - beta: SiTU gate ``beta`` (``activation_situ_beta``). - linear_beta: optional up-half ``linear_beta`` - (``activation_situ_linear_beta``); ``None`` keeps the up half - linear. - - Returns: - ``[num_tokens, d]`` tensor in ``x``'s dtype, numerically matching - the eager :class:`SituAndMul` reference. - """ - b, n = x.shape - - assert n % 2 == 0 - d = n // 2 - - o = torch.empty((b, d), dtype=x.dtype, device=x.device) - - def grid(meta: Mapping[str, int]) -> tuple[int, int]: - return (b, triton.cdiv(d, meta["BLOCK_SIZE"])) - - situ_and_mul_kernel[grid]( - o_ptr=o, - o_stride=o.stride(0), - x_ptr=x, - x_stride=x.stride(0), - d=d, - beta=float(beta), - linear_beta=float(linear_beta) if linear_beta is not None else 1.0, - BLOCK_SIZE=1024, - HAS_LINEAR_BETA=linear_beta is not None, - ) - - return o - - -@situ_and_mul.register_fake -def _(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor: - b, n = x.shape - - assert n % 2 == 0 - - return x.new_empty((b, n // 2)) - - -class NonSituActivation(nn.Module): - """SiLU/SwiGLU activation used as the non-SiTU mutation control. - - Splits the last dim into gate/up, applies SiLU to the gate, and - multiplies element-wise. Deliberately does NOT use the SiTU - ``beta * tanh(gate/beta) * sigmoid(gate)`` recipe. - """ - - def forward(self, x: torch.Tensor) -> torch.Tensor: - d = x.shape[-1] // 2 - gate = x[..., :d] - up = x[..., d:] - return torch.nn.functional.silu(gate) * up - - -class KimiK3MLP(nn.Module): - """K3 dense/shared-expert MLP module with TRT-LLM-style fused layout. - - Weight layout: - - * ``gate_up_proj``: ``nn.Linear(hidden_size, 2 * intermediate_size, bias=False)``. - Rows ``[:intermediate_size]`` correspond to HF's ``gate`` (KimiMLP.gate_proj - or KimiBlockSparseMLP.w1). Rows ``[intermediate_size:]`` correspond to - HF's ``up`` (KimiMLP.up_proj or KimiBlockSparseMLP.w3). - * ``down_proj``: ``nn.Linear(intermediate_size, hidden_size, bias=False)``. - Matches HF ``KimiMLP.down_proj`` or ``KimiBlockSparseMLP.w2``. - - Forward: ``down_proj( activation( gate_up_proj(x) ) )``. Default - ``activation`` is :class:`SituAndMul`; pass a different callable to - run mutation controls (e.g. :class:`NonSituActivation` for a - negative-control test). ``use_fused_activation=True`` routes CUDA - inputs through the fused Triton ``trtllm::situ_and_mul`` op instead - of the eager module (only valid with the default SiTU activation). - """ - - def __init__( - self, - *, - hidden_size: int, - intermediate_size: int, - situ_beta: float = 4.0, - situ_linear_beta: Optional[float] = 25.0, - activation: Optional[nn.Module] = None, - use_fused_activation: bool = False, - dtype: Optional[torch.dtype] = None, - device: Optional[torch.device] = None, - ) -> None: - super().__init__() - if use_fused_activation and activation is not None: - raise ValueError( - "use_fused_activation only fuses the default SiTU activation; " - "drop the custom activation module or the flag" - ) - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.use_fused_activation = use_fused_activation - - self.gate_up_proj = nn.Linear( - hidden_size, - 2 * intermediate_size, - bias=False, - dtype=dtype, - device=device, - ) - self.down_proj = nn.Linear( - intermediate_size, - hidden_size, - bias=False, - dtype=dtype, - device=device, - ) - self.activation = ( - activation - if activation is not None - else SituAndMul(beta=situ_beta, linear_beta=situ_linear_beta) - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - h1 = self.gate_up_proj(x) - if self.use_fused_activation and h1.is_cuda: - act = self.activation - h2 = torch.ops.trtllm.situ_and_mul( - h1.reshape(-1, h1.shape[-1]), act.beta, act.linear_beta - ).reshape(*h1.shape[:-1], self.intermediate_size) - else: - h2 = self.activation(h1) - return self.down_proj(h2) - - -class KimiK3RMSNorm(nn.Module): - """RMSNorm matching HF ``KimiRMSNorm`` semantics exactly. - - HF ``KimiRMSNorm.forward``:: - - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + eps) - return self.weight * hidden_states.to(input_dtype) - - ``self.weight`` in HF is initialised in the module's ambient dtype - (bf16 or fp32). Callers pin the weight dtype here too so byte-exact - parity holds regardless of the ambient dtype. - """ - - def __init__( - self, - hidden_size: int, - eps: float = 1e-6, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - ) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype, device=device)) - self.eps = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - # flashinfer's fused RMSNorm does the same fp32-accumulate - # normalization in one kernel, collapsing the eager - # pow/mean/rsqrt/mul/cast launch chain. Rounding differs by one final - # cast: flashinfer multiplies by ``weight`` in fp32 and casts once at - # the end, while the eager path casts the normalized value to the - # input dtype BEFORE the weight multiply — so outputs can differ by - # ~1 ulp and byte-exact HF parity requires the eager path. It is only - # valid for a CUDA fp16/bf16 input whose dtype matches the weight; - # CPU / fp32 parity paths, meta init, and the KIMI_K3_FUSED_RMSNORM=0 - # rollback keep the exact eager math below. - if ( - _FUSED_RMSNORM - and IS_FLASHINFER_AVAILABLE - and hidden_states.is_cuda - and hidden_states.dtype in (torch.float16, torch.bfloat16) - and self.weight.dtype == hidden_states.dtype - ): - from ...custom_ops import flashinfer_rmsnorm - - return flashinfer_rmsnorm(hidden_states.contiguous(), self.weight, self.eps) - input_dtype = hidden_states.dtype - h = hidden_states.to(torch.float32) - variance = h.pow(2).mean(-1, keepdim=True) - h = h * torch.rsqrt(variance + self.eps) - return self.weight * h.to(input_dtype) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py b/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py deleted file mode 100644 index 003175f2854f..000000000000 --- a/tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py +++ /dev/null @@ -1,281 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Kimi K3 MoE routing module. - -Structural mirror of HF ``KimiMoEGate`` (see -``model_config/modeling_kimi.py:710``). K3 routing inherits DeepSeek-V3's -``noaux_tc`` topology but pins the following K3 config choices (see -``configuration_kimi_k3.py``): - -* ``moe_router_activation_func = "sigmoid"`` — per-expert sigmoid, not softmax. -* ``e_score_correction_bias`` (per-expert) — bias added *only* for - top-k *selection*; the returned ``topk_weight`` samples the *raw* - sigmoid ``scores``, not the bias-adjusted ``scores_for_choice``. -* ``moe_renormalize = True`` — for ``top_k > 1``, divide ``topk_weight`` - by ``sum + 1e-20`` before scaling. -* ``routed_scaling_factor`` — final multiplicative scale. -* ``num_expert_group = 1``, ``topk_group = 1`` — K3 config disables the - grouped top-k branch used by DeepSeek. The gate still handles the - grouped branch when a caller flips those config knobs. - -Parameter names / shapes match HF ``KimiMoEGate`` so -:func:`copy_hf_moe_gate_weights` is identity name mapping. -""" - -from __future__ import annotations - -from typing import Any, Tuple - -import torch -from torch import nn - -from ..fused_moe.routing import DeepSeekV3MoeRoutingMethod, Deepseekv3RoutingImpl - - -class KimiK3MoEGate(nn.Module): - """K3 MoE routing — structural mirror of HF ``KimiMoEGate``. - - Positive path reproduces HF ``KimiMoEGate.forward`` at - ``modeling_kimi.py:747-803`` byte-identically under K3's - ``sigmoid`` scoring, top-k over the full expert set, raw sigmoid - weights, renormalization + scaling profile. - - Three mutation flags gate the negative controls required by AC6: - - * ``softmax_routing_mutation`` — softmax over experts instead of - per-expert sigmoid. - * ``biased_weights_mutation`` — gather ``topk_weight`` from the - bias-adjusted scores rather than the raw sigmoid scores. - * ``omit_renormalize_mutation`` — skip the renormalize step even - when the config asks for it. - """ - - def __init__( - self, - config: Any, - *, - softmax_routing_mutation: bool = False, - biased_weights_mutation: bool = False, - omit_renormalize_mutation: bool = False, - logits_gemm_dtype: torch.dtype | None = None, - device: torch.device | None = None, - ) -> None: - super().__init__() - self.config = config - self.top_k = config.num_experts_per_token - self.num_experts = config.num_experts - self.routed_scaling_factor = config.routed_scaling_factor - self.moe_router_activation_func = config.moe_router_activation_func - self.num_expert_group = getattr(config, "num_expert_group", 1) - self.topk_group = getattr(config, "topk_group", 1) - self.moe_renormalize = config.moe_renormalize - self.gating_dim = config.hidden_size - - assert self.moe_router_activation_func in ("sigmoid", "softmax"), ( - "K3 MoE gate supports sigmoid or softmax scoring only" - ) - - # Same parameter shapes / names as HF ``KimiMoEGate``. - # - # ``logits_gemm_dtype=torch.bfloat16`` stores the gate weight in - # bf16 and runs the logits GEMM as a single bf16xbf16 kernel with - # fp32 accumulate/output (``trtllm::dsv3_router_gemm_op``). The K3 - # checkpoint stores this weight in bf16, so the fp32 master was an - # exact upcast and bf16 storage is lossless; this removes the - # per-layer bf16->fp32 input cast + fp32 splitK-reduce that ran - # inside the decode CUDA graph (~5 us x 92 layers per step). - # Default ``None`` keeps the legacy fp32 GEMM (module parity tests). - weight_dtype = logits_gemm_dtype or torch.float32 - self.weight = nn.Parameter( - torch.empty((self.num_experts, self.gating_dim), dtype=weight_dtype, device=device) - ) - self.e_score_correction_bias = nn.Parameter(torch.empty(self.num_experts, device=device)) - - self.softmax_routing_mutation = softmax_routing_mutation - self.biased_weights_mutation = biased_weights_mutation - self.omit_renormalize_mutation = omit_renormalize_mutation - - # Fast path: the fused ``noaux_tc`` routing kernel computes exactly K3's - # production routing contract in one launch -- per-expert sigmoid, - # ``e_score_correction_bias`` added for *selection* only, top-k weights - # sampled from the raw sigmoid scores, renormalized by ``sum + 1e-20``, - # then scaled by ``routed_scaling_factor``. Route through the shared - # ``Deepseekv3RoutingImpl`` (same op DeepSeek-V3 uses) when the config is - # eligible and none of the parity-breaking mutation controls are active. - # The eager path below stays the reference for those controls, for - # softmax scoring, for ``moe_renormalize=False``, and for grouped / - # oversized configs the kernel does not support. - self._routing_impl = Deepseekv3RoutingImpl( - top_k=self.top_k, - n_group=self.num_expert_group, - topk_group=self.topk_group, - routed_scaling_factor=self.routed_scaling_factor, - is_fused=True, - ) - # Bounds mirror the n_group == 1 branch of - # ``Deepseekv3RoutingImpl.noaux_tc`` (num_experts <= 1024, top_k <= 32); - # staying inside them guarantees the fused kernel branch is taken (never - # the impl's own PyTorch fallback, whose grouped path differs from K3's). - self._use_fused_routing = ( - self.moe_router_activation_func == "sigmoid" - and self.num_expert_group == 1 - and self.moe_renormalize - and self.top_k > 1 - and self.num_experts <= 1024 - and self.top_k <= 32 - and not softmax_routing_mutation - and not biased_weights_mutation - and not omit_renormalize_mutation - ) - - def _score(self, logits: torch.Tensor) -> torch.Tensor: - if self.softmax_routing_mutation: - return logits.softmax(dim=1) - if self.moe_router_activation_func == "sigmoid": - return logits.sigmoid() - return logits.softmax(dim=1) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Routing logits ``[num_tokens, num_experts]``, fp32, pre-sigmoid. - - Used when the MoE block is hosted under ``ConfigurableMoE``: the - post-linear gate math (sigmoid, bias-for-selection, renormalize, - ``routed_scaling_factor``) runs inside the wrapper's routing method - per chunk; only the gate GEMM stays here, keeping the checkpoint - parameter mapping identity. - """ - hidden_2d = hidden_states.reshape(-1, self.gating_dim) - if self.weight.dtype == torch.bfloat16 and hidden_2d.dtype == torch.bfloat16: - # Single bf16xbf16 -> fp32 GEMM (fp32 accumulate); no input - # upcast kernel, no fp32 splitK-reduce. K3's 896 experts miss - # the op's specialized 256-expert kernels and take its cublas - # path, which is the point here (one fused kernel). - return torch.ops.trtllm.dsv3_router_gemm_op( - hidden_2d.contiguous(), - self.weight.t(), - bias=None, - out_dtype=torch.float32, - ) - return torch.nn.functional.linear( - hidden_2d.type(torch.float32), - self.weight.type(torch.float32), - None, - ) - - @property - def routing_method(self) -> DeepSeekV3MoeRoutingMethod: - """Return the shared DeepSeekV3 router used by ConfigurableMoE.""" - if self.moe_router_activation_func != "sigmoid": - raise ValueError("Kimi K3 ConfigurableMoE routing requires sigmoid scores.") - if not self.moe_renormalize: - raise ValueError( - "Kimi K3 ConfigurableMoE routing requires top-k weight renormalization." - ) - if ( - self.softmax_routing_mutation - or self.biased_weights_mutation - or self.omit_renormalize_mutation - ): - raise ValueError( - "Kimi K3 routing mutation flags are reference-test controls " - "and cannot be used by ConfigurableMoE." - ) - return DeepSeekV3MoeRoutingMethod( - top_k=self.top_k, - n_group=self.num_expert_group, - topk_group=self.topk_group, - routed_scaling_factor=self.routed_scaling_factor, - callable_e_score_correction_bias=lambda: self.e_score_correction_bias, - is_fused=True, - ) - - def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - logits = self.compute_logits(hidden_states) - # ``compute_logits`` flattens to [num_tokens, num_experts]; derive - # the token count from it so any input rank works. - num_tokens = logits.shape[0] - - # ``trtllm::noaux_tc_op`` is a CUDA-only custom op; CPU inputs - # (reference / parity tests) fall through to the eager path below, - # which stays the routing-contract reference on every device. - if self._use_fused_routing and logits.is_cuda: - # One fused kernel replaces the sigmoid -> (+bias) -> top-k -> - # gather -> renormalize -> scale chain below. ``noaux_tc`` returns - # (weights, indices); return the eager dtype contract -- int64 - # indices (as ``torch.topk`` yields) and fp32 weights -- so every - # downstream consumer is byte-for-byte unaffected by the swap. - topk_weight, topk_idx = self._routing_impl.noaux_tc( - logits, self.e_score_correction_bias.float() - ) - return topk_idx.to(torch.int64), topk_weight.to(torch.float32) - - scores = self._score(logits) - scores = scores.view(num_tokens, -1) - - # Bias is applied for *selection*, not for the returned weight. - scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) - - if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: - group_scores = ( - scores_for_choice.view(num_tokens, self.num_expert_group, -1) - .topk(2, dim=-1)[0] - .sum(dim=-1) - ) - group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] - group_mask = torch.zeros_like(group_scores) - group_mask.scatter_(1, group_idx, 1) - score_mask = ( - group_mask.unsqueeze(-1) - .expand( - num_tokens, - self.num_expert_group, - self.num_experts // self.num_expert_group, - ) - .reshape(num_tokens, -1) - ) - tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) - else: - tmp_scores = scores_for_choice - - _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False) - - # Positive contract: gather from raw ``scores`` (not bias-adjusted). - weight_source = scores_for_choice if self.biased_weights_mutation else scores - topk_weight = weight_source.gather(1, topk_idx) - - if self.top_k > 1 and self.moe_renormalize and not self.omit_renormalize_mutation: - denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20 - topk_weight = topk_weight / denominator - - topk_weight = topk_weight * self.routed_scaling_factor - return topk_idx, topk_weight - - -def copy_hf_moe_gate_weights( - hf: nn.Module, - k3: KimiK3MoEGate, -) -> dict[str, tuple[tuple[int, ...], str]]: - """Copy parameters from HF ``KimiMoEGate`` into ``k3``. - - Identity name mapping (``weight`` + ``e_score_correction_bias``). - Returns a ``{name: (shape, dtype)}`` provenance dict. - """ - src_params = dict(hf.named_parameters()) - dst_params = dict(k3.named_parameters()) - missing_on_k3 = sorted(set(src_params) - set(dst_params)) - missing_on_hf = sorted(set(dst_params) - set(src_params)) - if missing_on_k3: - raise KeyError(f"copy_hf_moe_gate_weights: HF params missing on K3: {missing_on_k3}") - if missing_on_hf: - raise KeyError(f"copy_hf_moe_gate_weights: K3 params missing on HF: {missing_on_hf}") - provenance = {} - for name, src in src_params.items(): - dst = dst_params[name] - if src.shape != dst.shape: - raise ValueError( - f"shape mismatch for {name}: HF {tuple(src.shape)} vs K3 {tuple(dst.shape)}" - ) - with torch.no_grad(): - dst.data.copy_(src.data.to(dtype=dst.dtype, device=dst.device)) - provenance[name] = (tuple(src.shape), str(src.dtype)) - return provenance diff --git a/tensorrt_llm/_torch/modules/situ.py b/tensorrt_llm/_torch/modules/situ.py new file mode 100644 index 000000000000..e40e60e34337 --- /dev/null +++ b/tensorrt_llm/_torch/modules/situ.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SiTU gated activation and its fused Triton implementation.""" + +from __future__ import annotations + +from typing import Mapping, Optional + +import torch +import triton # type: ignore[import] +import triton.language as tl # type: ignore[import] +import triton.language.extra.libdevice as tldevice # type: ignore[import] +from torch import nn + + +class SituAndMul(nn.Module): + """SiTU activation with gate/up multiplicative gating. + + Runs the math in fp32 for numerical stability, then casts back to the + input dtype. CUDA inputs optionally use the fused Triton implementation; + CPU and meta inputs keep the eager reference path. + """ + + def __init__( + self, + *, + beta: float = 1.0, + linear_beta: Optional[float] = None, + use_fused_activation: bool = False, + ) -> None: + super().__init__() + self.beta = beta + self.linear_beta = linear_beta + self.use_fused_activation = use_fused_activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.use_fused_activation and x.is_cuda: + return torch.ops.trtllm.situ_and_mul( + x.reshape(-1, x.shape[-1]), self.beta, self.linear_beta + ).reshape(*x.shape[:-1], x.shape[-1] // 2) + + d = x.shape[-1] // 2 + gate = x[..., :d].to(torch.float32) + up = x[..., d:].to(torch.float32) + situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate) + if self.linear_beta is not None: + up = self.linear_beta * torch.tanh(up / self.linear_beta) + return (situ_a * up).to(x.dtype) + + +@triton.jit +def situ_and_mul_kernel( + o_ptr, + o_stride, + x_ptr, + x_stride, + d, + beta, + linear_beta, + BLOCK_SIZE: tl.constexpr, + HAS_LINEAR_BETA: tl.constexpr, +) -> None: + """Fused :class:`SituAndMul` on a packed ``[gate | up]`` row layout.""" + i = tl.program_id(axis=0).to(tl.int64) + j = tl.program_id(axis=1) + + o_row_ptr = o_ptr + o_stride * i + x_row_ptr = x_ptr + x_stride * i + + offsets = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < d + + gate = tl.load(x_row_ptr + offsets, mask=mask).to(tl.float32) + up = tl.load(x_row_ptr + offsets + d, mask=mask).to(tl.float32) + + situ_a = beta * tldevice.tanh(gate / beta) * tl.sigmoid(gate) + if HAS_LINEAR_BETA: + up = linear_beta * tldevice.tanh(up / linear_beta) + result = situ_a * up + + tl.store(o_row_ptr + offsets, result, mask=mask) + + +@torch.library.custom_op("trtllm::situ_and_mul", mutates_args=()) +def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor: + """Run fused SiTU on a packed ``[num_tokens, 2 * d]`` input.""" + b, n = x.shape + + assert n % 2 == 0 + d = n // 2 + output = torch.empty((b, d), dtype=x.dtype, device=x.device) + + def grid(meta: Mapping[str, int]) -> tuple[int, int]: + return (b, triton.cdiv(d, meta["BLOCK_SIZE"])) + + situ_and_mul_kernel[grid]( + o_ptr=output, + o_stride=output.stride(0), + x_ptr=x, + x_stride=x.stride(0), + d=d, + beta=float(beta), + linear_beta=float(linear_beta) if linear_beta is not None else 1.0, + BLOCK_SIZE=1024, + HAS_LINEAR_BETA=linear_beta is not None, + ) + return output + + +@situ_and_mul.register_fake +def _(x: torch.Tensor, beta: float, linear_beta: Optional[float] = None) -> torch.Tensor: + b, n = x.shape + assert n % 2 == 0 + return x.new_empty((b, n // 2)) diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py index e6873815f06b..f21b81328722 100644 --- a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py +++ b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_moe_kernels.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Native TRTLLM-Gen SiTU MoE dispatch for the in-tree Kimi K3 sparse MoE block. +"""Test-only native SiTU dispatch for the Kimi K3 sparse MoE reference. The K3 MoE block has two mutually exclusive kernel paths: @@ -127,7 +127,7 @@ def pack_routed_expert_weights( """Pad + shuffle checkpoint MXFP4 expert weights into the TRTLLM-Gen layout. Inputs are the per-expert MXFP4 tensors as stored by - :class:`KimiK3RoutedExpertBank` (HF layout, group_size=32): + the test reference's routed expert bank (HF layout, group_size=32): * ``w1_packed``/``w3_packed``: ``uint8 [E, I, H // 2]`` (w1 = gate, w3 = up) * ``w1_scales``/``w3_scales``: ``uint8 [E, I, H // 32]`` (E8M0 biased exponents) diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py index 7d91294462a1..b7e12491a234 100644 --- a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py +++ b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/_mxfp4.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""MXFP4 (E2M1 + E8M0) group-scaled quantization utilities. +"""Test-only MXFP4 (E2M1 + E8M0) quantization utilities. Kimi K3's routed expert Linear weights are stored in the ``mxfp4-pack-quantized`` format with ``group_size=32`` (see @@ -149,15 +149,3 @@ def dequantize_last_dim_mxfp4( codes_grouped = codes_flat.reshape(*lead, num_groups, group_size) deq_grouped = _dequantize_group_e2m1(codes_grouped, scales_u8) return deq_grouped.reshape(*lead, n) - - -def canonical_mxfp4_fp32(x: torch.Tensor, group_size: int = DEFAULT_GROUP_SIZE) -> torch.Tensor: - """Round-trip ``x`` (fp32) through MXFP4 pack and unpack. - - Convenience for tests: yields the fp32 value that a stored MXFP4 - weight actually decodes to. Callers use this to initialize a - reference (unquantized) module with the same values a K3 MXFP4 - weight produces, giving byte-exact parity. - """ - packed, scales = quantize_last_dim_mxfp4(x, group_size=group_size) - return dequantize_last_dim_mxfp4(packed, scales, group_size=group_size) diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py new file mode 100644 index 000000000000..f316b3a16158 --- /dev/null +++ b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_mlp_test_utils.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Test-only Kimi K3 MLP reference modules.""" + +import torch +from torch import nn + +from tensorrt_llm._torch.modules.situ import SituAndMul + + +class NonSituActivation(nn.Module): + """SiLU/SwiGLU activation used as a non-SiTU mutation control.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + d = x.shape[-1] // 2 + gate = x[..., :d] + up = x[..., d:] + return torch.nn.functional.silu(gate) * up + + +class KimiK3MLP(nn.Module): + """K3 MLP reference with a fused ``gate_up_proj`` weight layout.""" + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + situ_beta: float = 4.0, + situ_linear_beta: float | None = 25.0, + activation: nn.Module | None = None, + use_fused_activation: bool = False, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + ) -> None: + super().__init__() + if use_fused_activation and activation is not None: + raise ValueError( + "use_fused_activation only fuses the default SiTU activation; " + "drop the custom activation module or the flag" + ) + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.use_fused_activation = use_fused_activation + + self.gate_up_proj = nn.Linear( + hidden_size, + 2 * intermediate_size, + bias=False, + dtype=dtype, + device=device, + ) + self.down_proj = nn.Linear( + intermediate_size, + hidden_size, + bias=False, + dtype=dtype, + device=device, + ) + self.activation = ( + activation + if activation is not None + else SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=use_fused_activation, + ) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.activation(self.gate_up_proj(x))) diff --git a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py index a493f847fccc..1e3f58e5938c 100644 --- a/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py +++ b/tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/kimi_k3_moe_block.py @@ -7,7 +7,7 @@ ``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:`KimiK3ReferenceMoEGate` for eager reference routing. * :class:`KimiK3RoutedExpertBank` — per-expert MXFP4-packed ``w1 / w2 / w3`` linear weights (group_size=32), dequantized on the fly during the Python fallback path. @@ -59,15 +59,53 @@ dequantize_last_dim_mxfp4, quantize_last_dim_mxfp4, ) +from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP, NonSituActivation 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 +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate, KimiK3RMSNorm +from tensorrt_llm._torch.modules.situ import SituAndMul + + +class KimiK3ReferenceMoEGate(KimiK3MoEGate): + """Eager HF-compatible Kimi K3 routing reference.""" + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + logits = self.compute_logits(hidden_states) + num_tokens = logits.shape[0] + if self.moe_router_activation_func == "sigmoid": + scores = logits.sigmoid() + else: + scores = logits.softmax(dim=1) + scores = scores.view(num_tokens, -1) + + scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0) + if self.num_expert_group > 1 and self.num_expert_group > self.topk_group: + group_scores = ( + scores_for_choice.view(num_tokens, self.num_expert_group, -1) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) + group_indices = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_indices, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand( + num_tokens, + self.num_expert_group, + self.num_experts // self.num_expert_group, + ) + .reshape(num_tokens, -1) + ) + scores_for_selection = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf")) + else: + scores_for_selection = scores_for_choice + + topk_indices = torch.topk(scores_for_selection, k=self.top_k, dim=-1, sorted=False)[1] + topk_weights = scores.gather(1, topk_indices) + if self.top_k > 1 and self.moe_renormalize: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + return topk_indices, topk_weights * self.routed_scaling_factor class KimiK3RoutedExpertBank(nn.Module): @@ -315,7 +353,7 @@ def __init__( else: routed_activation = SituAndMul(beta=situ_beta, linear_beta=situ_linear_beta) - self.gate = KimiK3MoEGate(config, device=device) + self.gate = KimiK3ReferenceMoEGate(config, device=device) # Routed expert storage — the MXFP4 bank is always the checkpoint # quantization source of truth. The fused path derives its packed diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index 593c46e755c0..e0fd9da3c2f4 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -2,11 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """Kimi K3 fused gate_up MLP tests. -The runtime dense / shared-expert MLP (``KimiK3MLP``) runs a single fused -``gate_up_proj`` GEMM whose weight is the row-concat of the HF checkpoint's -separate ``gate_proj`` / ``up_proj`` tensors (the same concat -``KimiLinearForCausalLM.load_weights`` performs). These tests check the -fused module against an unfused reference built from the split weights: +The runtime dense / shared-expert ``GatedMLP`` runs a single fused +``gate_up_proj`` GEMM with K3's SiTU activation. ``KimiK3MLP`` remains the +compact reference for the same fused weight layout. These tests check that +layout against an unfused reference built from the HF checkpoint's split +``gate_proj`` / ``up_proj`` weights: * fused ``gate_up_proj`` output matches ``two GEMMs + torch.cat`` + eager ``SituAndMul`` + ``down_proj`` for a decode-shaped (1 token) and @@ -14,13 +14,18 @@ ``None`` (the two activation code paths); * the row-concat convention is required: swapping the halves breaks the numerics (mutation control). +* direct MoE-TP selects combined reduction for the shared and routed partials. """ +from types import SimpleNamespace + import pytest import torch +from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP from torch import nn -from tensorrt_llm._torch.modules.kimi_k3_moe._mlp import KimiK3MLP, SituAndMul +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.situ import SituAndMul requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") @@ -68,6 +73,26 @@ def _make_pair(hidden_size, intermediate_size, situ_beta, situ_linear_beta, devi return fused, ref +def _runtime_config() -> SimpleNamespace: + return SimpleNamespace( + hidden_size=512, + num_experts=8, + num_experts_per_token=2, + moe_intermediate_size=256, + num_shared_experts=2, + routed_expert_hidden_size=256, + latent_moe_use_norm=True, + rms_norm_eps=1e-5, + activation_situ_beta=4.0, + activation_situ_linear_beta=25.0, + moe_renormalize=True, + moe_router_activation_func="sigmoid", + routed_scaling_factor=1.0, + num_expert_group=1, + topk_group=1, + ) + + @requires_cuda @pytest.mark.parametrize( "situ_beta,situ_linear_beta", @@ -115,3 +140,196 @@ def test_gate_up_half_swap_mutation_breaks_accuracy(): x = torch.randn(64, hidden_size, dtype=torch.bfloat16, device=device) * 0.5 with pytest.raises(AssertionError): torch.testing.assert_close(fused(x), ref(x), rtol=1.6e-2, atol=1e-3) + + +@requires_cuda +@pytest.mark.parametrize( + "situ_beta,situ_linear_beta", + [(4.0, 25.0), (1.0, None)], + ids=["default", "no_linear_beta"], +) +def test_gated_mlp_supports_fused_situ(situ_beta, situ_linear_beta): + """The shared-expert replacement preserves K3 MLP numerics.""" + device = torch.device("cuda") + hidden_size, intermediate_size = 512, 384 + torch.manual_seed(17) + reference = KimiK3MLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + use_fused_activation=True, + dtype=torch.bfloat16, + device=device, + ) + gated = GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + activation=SituAndMul( + beta=situ_beta, + linear_beta=situ_linear_beta, + use_fused_activation=True, + ), + dtype=torch.bfloat16, + reduce_output=True, + ).to(device) + + with torch.no_grad(): + for projection in (reference.gate_up_proj, reference.down_proj): + projection.weight.copy_(torch.randn_like(projection.weight, dtype=torch.float32) * 0.05) + gated.gate_up_proj.weight.copy_(reference.gate_up_proj.weight) + gated.down_proj.weight.copy_(reference.down_proj.weight) + + x = torch.randn(64, hidden_size, dtype=torch.bfloat16, device=device) * 0.5 + torch.testing.assert_close(gated(x), reference(x), rtol=1.6e-2, atol=1e-3) + + +@pytest.mark.parametrize( + "attention_dp,tp_size,has_routed_comm,expected_shared_tp,expected_shared_reduce,expected_combined_ar", + [ + (True, 2, True, 1, False, False), + (False, 1, False, 1, False, False), + (False, 2, False, 2, False, True), + (False, 2, True, 2, True, False), + ], + ids=["attention_dp", "single_rank", "combined_tp", "routed_comm"], +) +def test_kimi_k3_shared_expert_reduction_mode( + monkeypatch, + attention_dp, + tp_size, + has_routed_comm, + expected_shared_tp, + expected_shared_reduce, + expected_combined_ar, +): + """Each parallel mode selects exactly the required shared reduction.""" + from tensorrt_llm._torch import distributed + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models import modeling_kimi_linear + from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantConfig + + class _FakeAllReduce(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + fake_moe = ConfigurableMoE.__new__(ConfigurableMoE) + torch.nn.Module.__init__(fake_moe) + fake_moe.backend = SimpleNamespace(initial_local_expert_ids=[0, 1, 2, 3]) + fake_moe.comm = object() if has_routed_comm else None + fake_moe.layer_load_balancer = None + fake_moe.all_reduce = _FakeAllReduce() if not attention_dp and tp_size > 1 else None + + monkeypatch.setattr(modeling_kimi_linear, "create_moe", lambda **_: fake_moe) + monkeypatch.setattr(modeling_kimi_linear, "AllReduce", _FakeAllReduce) + monkeypatch.setattr(distributed, "AllReduce", _FakeAllReduce) + monkeypatch.setattr(torch.cuda, "Event", lambda: object()) + + mapping = Mapping( + world_size=tp_size, + rank=0, + tp_size=tp_size, + enable_attention_dp=attention_dp, + ) + model_config = ModelConfig(mapping=mapping, quant_config=QuantConfig()) + config = _runtime_config() + runtime = modeling_kimi_linear.KimiK3MoERuntime(model_config, config, layer_idx=1) + + shared = runtime.shared_experts + assert isinstance(shared, GatedMLP) + assert shared.gate_up_proj.tp_size == expected_shared_tp + assert shared.down_proj.tp_size == expected_shared_tp + assert shared.down_proj.reduce_output is expected_shared_reduce + assert runtime._use_combined_all_reduce is expected_combined_ar + local_intermediate = 512 // expected_shared_tp + assert shared.gate_up_proj.weight.shape == (2 * local_intermediate, config.hidden_size) + assert shared.down_proj.weight.shape == (config.hidden_size, local_intermediate) + + +@pytest.mark.parametrize( + "attention_dp,tp_size,rank,gpus_per_node,intermediate_size,expected_tp_size,expected_tp_rank", + [ + (True, 8, 7, 8, 516, 1, 0), + (False, 8, 7, 8, 512, 8, 7), + (False, 8, 7, 8, 516, 4, 3), + (False, 8, 7, 8, 515, 1, 0), + (False, 16, 15, 8, 512, 8, 7), + ], + ids=["attention_dp", "full_tp", "gcd_subgroup", "replicated", "single_node_cap"], +) +def test_kimi_k3_dense_layer_uses_gated_mlp( + monkeypatch, + attention_dp, + tp_size, + rank, + gpus_per_node, + intermediate_size, + expected_tp_size, + expected_tp_rank, +): + """The first dense layer selects a block-aligned, node-local MLP TP group.""" + from tensorrt_llm._torch import distributed + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models import modeling_kimi_linear + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantConfig + + class _IdentityAttention(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, hidden_states, attn_metadata): + return hidden_states + + class _IdentityAllReduce(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + def forward(self, hidden_states, *args, **kwargs): + return hidden_states + + monkeypatch.setattr(modeling_kimi_linear, "KimiKDARuntime", _IdentityAttention) + monkeypatch.setattr(distributed, "AllReduce", _IdentityAllReduce) + + mapping = Mapping( + world_size=tp_size, + rank=rank, + gpus_per_node=gpus_per_node, + tp_size=tp_size, + enable_attention_dp=attention_dp, + ) + model_config = ModelConfig(mapping=mapping, quant_config=QuantConfig()) + config = SimpleNamespace( + hidden_size=512, + intermediate_size=intermediate_size, + num_experts=8, + first_k_dense_replace=1, + moe_layer_freq=1, + linear_attn_config={"kda_layers": [1], "full_attn_layers": []}, + rms_norm_eps=1e-5, + attn_res_block_size=1, + activation_situ_beta=4.0, + activation_situ_linear_beta=25.0, + ) + layer = modeling_kimi_linear.KimiLinearDecoderLayer(model_config, config, layer_idx=0) + + assert not layer.is_moe + assert isinstance(layer.mlp, GatedMLP) + assert layer.mlp_tp_size == expected_tp_size + assert layer.mlp.gate_up_proj.tp_size == expected_tp_size + assert layer.mlp.gate_up_proj.tp_rank == expected_tp_rank + assert layer.mlp.down_proj.tp_size == expected_tp_size + assert layer.mlp.down_proj.tp_rank == expected_tp_rank + assert layer.mlp.down_proj.reduce_output is (expected_tp_size > 1) + local_intermediate = config.intermediate_size // expected_tp_size + assert layer.mlp.gate_up_proj.weight.shape == ( + 2 * local_intermediate, + config.hidden_size, + ) + assert layer.mlp.down_proj.weight.shape == ( + config.hidden_size, + local_intermediate, + ) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py index 7027a118cab3..a887c98723d8 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py @@ -1,29 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Parity tests for the Kimi K3 MoE gate fused-routing fast path. - -``KimiK3MoEGate.forward`` routes eligible configs (per-expert sigmoid -scoring, ``num_expert_group == 1``, renormalize on, no mutation controls, -within the kernel's supported bounds) through the fused -``torch.ops.trtllm.noaux_tc_op`` kernel via ``Deepseekv3RoutingImpl``. - -These tests assert: - -* the fused path is numerically equivalent to the eager - ``sigmoid -> +bias select -> top-k -> gather raw scores -> renormalize - (sum + 1e-20) -> scale`` reference it replaces (same experts selected, - same weights), and preserves the eager dtype contract (int64 indices, - fp32 weights) every downstream consumer relies on; and -* ineligible configs (softmax scoring, grouped routing, renormalize off, - any mutation control) keep the eager reference path. -""" +"""Parity tests for the production Kimi K3 MoE routing method.""" import dataclasses import pytest import torch +from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_moe_block import KimiK3ReferenceMoEGate -from tensorrt_llm._torch.modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate @dataclasses.dataclass @@ -55,34 +40,25 @@ def _dense_weights( @pytest.mark.skipif(not torch.cuda.is_available(), reason="noaux_tc_op is a CUDA custom op") -def test_fused_routing_matches_eager(): +def test_fused_routing_matches_eager_reference(): torch.manual_seed(0) cfg = _GateConfig() gate = KimiK3MoEGate(cfg).cuda() + reference_gate = KimiK3ReferenceMoEGate(cfg).cuda() with torch.no_grad(): gate.weight.normal_(0.0, 0.05) gate.e_score_correction_bias.normal_(0.0, 0.1) - - # The stock K3 config takes the fused fast path. - assert gate._use_fused_routing is True + reference_gate.load_state_dict(gate.state_dict()) num_tokens = 17 hidden = torch.randn(1, num_tokens, cfg.hidden_size, device="cuda") - idx_fused, wt_fused = gate(hidden) + idx_eager, wt_eager = reference_gate(hidden) + idx_fused, wt_fused = gate.routing_method.apply(gate.compute_logits(hidden)) - # Dtype/shape contract preserved for downstream consumers: the python - # fallback's ``scatter_`` needs int64 indices; the weighted sum consumes - # fp32 weights. - assert idx_fused.dtype == torch.int64 - assert wt_fused.dtype == torch.float32 assert idx_fused.shape == (num_tokens, cfg.num_experts_per_token) assert wt_fused.shape == (num_tokens, cfg.num_experts_per_token) - # Force the eager reference path on the same gate/weights/input. - gate._use_fused_routing = False - idx_eager, wt_eager = gate(hidden) - # Same experts selected per token (top-k is unsorted, so compare sets). sel_fused = torch.sort(idx_fused, dim=-1).values sel_eager = torch.sort(idx_eager.to(torch.int64), dim=-1).values @@ -98,30 +74,12 @@ def test_fused_routing_matches_eager(): @pytest.mark.parametrize( - "kwargs", - [ - dict(softmax_routing_mutation=True), - dict(biased_weights_mutation=True), - dict(omit_renormalize_mutation=True), - ], -) -def test_mutation_controls_disable_fused_routing(kwargs): - # The mutation controls change the routing math, so they must fall back - # to the eager reference rather than the fused kernel. - gate = KimiK3MoEGate(_GateConfig(), **kwargs) - assert gate._use_fused_routing is False - - -@pytest.mark.parametrize( - "cfg", + "cfg,error", [ - _GateConfig(moe_router_activation_func="softmax"), - _GateConfig(num_expert_group=4, topk_group=2), - _GateConfig(moe_renormalize=False), - _GateConfig(num_experts_per_token=1), + (_GateConfig(moe_router_activation_func="softmax"), "requires sigmoid"), + (_GateConfig(moe_renormalize=False), "requires top-k weight renormalization"), ], ) -def test_ineligible_configs_disable_fused_routing(cfg): - # softmax scoring, grouped routing, renormalize off, and top_k == 1 all - # diverge from the fused kernel's fixed contract -> eager path. - assert KimiK3MoEGate(cfg)._use_fused_routing is False +def test_invalid_config_rejected_by_production_routing(cfg, error): + with pytest.raises(ValueError, match=error): + _ = KimiK3MoEGate(cfg).routing_method diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py index bcf91178a8c8..f5bb822a6f0c 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py @@ -18,8 +18,9 @@ import pytest import torch +from _torch.modules.moe.kimi_k3_ref_moe.kimi_k3_mlp_test_utils import KimiK3MLP -from tensorrt_llm._torch.modules.kimi_k3_moe._mlp import KimiK3MLP, SituAndMul +from tensorrt_llm._torch.modules.situ import SituAndMul requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") 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 733d7eedb306..1e0648ccca88 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 @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Kimi K3 native TRTLLM-Gen SiTU MoE tests. -Covers the acceptance criteria of the SiTU cubin integration plan -(`tensorrt_llm/_torch/modules/kimi_k3_moe/SITU_CUBIN_INTEGRATION_PLAN.md`): +Covers the SiTU cubin integration behavior: * runner-local ActType numeric stability (SwiGlu/Relu2/Silu unchanged, SiTu appended); @@ -41,13 +40,13 @@ from tensorrt_llm._torch.models.modeling_kimi_linear import ( _K3_MOE_EP_ENV, _K3_MOE_TP_ENV, + KimiK3MoEGate, KimiK3MoERuntime, ) from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory from tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_deepgemm import ( _MEGA_MOE_SYMM_BUFFER_CACHE, ) -from tensorrt_llm._torch.modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate from tensorrt_llm._torch.utils import ActType_TrtllmGen from tensorrt_llm._utils import get_free_port from tensorrt_llm.mapping import Mapping @@ -180,12 +179,19 @@ def test_kimi_gate_reuses_deepseek_v3_routing(): gate.e_score_correction_bias.normal_(std=0.05) hidden_states = torch.randn(2, 7, config.hidden_size) - expected_ids, expected_weights = gate(hidden_states) + logits = gate.compute_logits(hidden_states) + scores = logits.sigmoid() + scores_for_choice = scores + gate.e_score_correction_bias.unsqueeze(0) + expected_ids = scores_for_choice.topk(gate.top_k, dim=-1, sorted=False).indices + expected_weights = scores.gather(1, expected_ids) + expected_weights /= expected_weights.sum(dim=-1, keepdim=True) + expected_weights *= gate.routed_scaling_factor + routing_method = gate.routing_method # Exercise the portable PyTorch short path; the production path keeps # is_fused=True and uses the same routing contract. routing_method.routing_impl.is_fused = False - actual_ids, actual_weights = routing_method.apply(gate.compute_logits(hidden_states)) + actual_ids, actual_weights = routing_method.apply(logits) expected_order = expected_ids.argsort(dim=-1) actual_order = actual_ids.argsort(dim=-1) From 5d1e88f0b33298b2b59fdf2e59ab3afe6135ff9f Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 13 Aug 2026 21:10:13 +0800 Subject: [PATCH 2/9] fix(torch): address Kimi K3 MLP review Signed-off-by: Jiagan Cheng --- .../_torch/models/modeling_kimi_linear.py | 43 +++++++-------- tensorrt_llm/_torch/modules/situ.py | 2 + .../test_lists/test-db/l0_b200.yml | 1 + .../_torch/modules/moe/test_kimi_k3_mlp.py | 53 ++++++++++++++++++- .../modules/moe/test_kimi_k3_situ_and_mul.py | 10 ++++ 5 files changed, 83 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index cf1065e055d9..0071239d0393 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -921,9 +921,14 @@ def __init__( # all-reduce. Communication-backed MoE already returns a complete # routed result, so GatedMLP must reduce only its shared partial. use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1 - self._use_combined_all_reduce = ( - use_shared_tp and routed_comm is None and self.routed_experts.all_reduce is not None - ) + needs_routed_all_reduce = use_shared_tp and routed_comm is None + routed_all_reduce = self.routed_experts.all_reduce + if needs_routed_all_reduce and routed_all_reduce is None: + raise RuntimeError( + "Kimi K3 direct MoE tensor parallelism requires the " + "ConfigurableMoE all-reduce even when reduce_results=False." + ) + self._use_combined_all_reduce = needs_routed_all_reduce self.shared_experts = GatedMLP( hidden_size=cfg.hidden_size, intermediate_size=shared_intermediate, @@ -959,13 +964,12 @@ def __init__( hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype ) - def _routed_up_projection(self, hidden_states: torch.Tensor) -> torch.Tensor: - if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance( - self.routed_expert_up_proj, nn.Linear - ): - return self.routed_expert_up_proj(hidden_states) + @staticmethod + def _routed_projection(hidden_states: torch.Tensor, projection: nn.Module) -> torch.Tensor: + if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance(projection, nn.Linear): + return projection(hidden_states) return torch.ops.trtllm.dsv3_fused_a_gemm_op( - hidden_states, self.routed_expert_up_proj.weight.t(), None, None + hidden_states, projection.weight.t(), None, None ) @staticmethod @@ -1077,7 +1081,7 @@ def forward(self, hidden_states: torch.Tensor, all_rank_num_tokens=None) -> torc """``hidden_states``: ``[num_tokens, hidden_size]`` bf16.""" identity = hidden_states router_logits = self.gate.compute_logits(hidden_states) - moe_all_reduce = self.routed_experts.all_reduce + moe_all_reduce = self.routed_experts.all_reduce if self._use_combined_all_reduce else None def _routed_output(): # Latent down/up projections via the min-latency fused GEMM op: @@ -1090,14 +1094,7 @@ def _routed_output(): # replaced the projection module, call it directly: its weight is # an e4m3 buffer the bf16 dsv3 op must not read, and its forward # is already a single fused GEMM (fp8_swap_ab_gemm). - if _K3_DISABLE_MIN_LATENCY_LATENT_PROJ or not isinstance( - self.routed_expert_down_proj, nn.Linear - ): - routed_in = self.routed_expert_down_proj(hidden_states) - else: - routed_in = torch.ops.trtllm.dsv3_fused_a_gemm_op( - hidden_states, self.routed_expert_down_proj.weight.t(), None, None - ) + routed_in = self._routed_projection(hidden_states, self.routed_expert_down_proj) y = self.routed_experts( routed_in, router_logits, @@ -1105,13 +1102,9 @@ def _routed_output(): ) if self._use_combined_all_reduce: return y - # EP partial latent sums must be completed before the nonlinear - # latent norm. Communication-backed paths combine internally; - # direct TP paths use the wrapper's AllReduce explicitly. - if self.routed_experts.comm is None and moe_all_reduce is not None: - y = moe_all_reduce(y) + # Communication-backed paths return a complete routed result. y = self.routed_expert_norm(y) - return self._routed_up_projection(y) + return self._routed_projection(y, self.routed_expert_up_proj) # Shared experts depend only on the block input, so overlap their GEMMs # with the routed dispatch/expert/combine chain. Multi-stream engages @@ -1136,7 +1129,7 @@ def _routed_output(): # The column split is a strided view; FlashInfer RMSNorm expects # a dense last dimension. routed_latent = self.routed_expert_norm(routed_latent.contiguous()) - routed_out = self._routed_up_projection(routed_latent) + routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) return routed_out + shared_out diff --git a/tensorrt_llm/_torch/modules/situ.py b/tensorrt_llm/_torch/modules/situ.py index e40e60e34337..980893271bfe 100644 --- a/tensorrt_llm/_torch/modules/situ.py +++ b/tensorrt_llm/_torch/modules/situ.py @@ -87,6 +87,8 @@ def situ_and_mul(x: torch.Tensor, beta: float, linear_beta: Optional[float] = No b, n = x.shape assert n % 2 == 0 + if x.stride(1) != 1: + raise ValueError("situ_and_mul requires a contiguous last dimension") d = n // 2 output = torch.empty((b, d), dtype=x.dtype, device=x.device) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 9f9b2348103b..8a6ecb3ce6e8 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -106,6 +106,7 @@ l0_b200: # GPU KDA disagg transfer + peer-validation (cpu_only cases skipped by the # stage's "not cpu_only" markexpr and run on the CPU-Generic stage instead). - unittest/disaggregated/test_kda_mamba_transfer.py + - unittest/_torch/modules/moe/test_kimi_k3_mlp.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) --------------- diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index e0fd9da3c2f4..57569bdec7c1 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -185,6 +185,57 @@ def test_gated_mlp_supports_fused_situ(situ_beta, situ_linear_beta): torch.testing.assert_close(gated(x), reference(x), rtol=1.6e-2, atol=1e-3) +@requires_cuda +def test_create_moe_keeps_all_reduce_for_external_tp_reduction(): + """``reduce_results=False`` keeps the wrapper's direct-TP collective.""" + from transformers.configuration_utils import PretrainedConfig + + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate + from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE, create_moe + from tensorrt_llm._torch.utils import ActType_TrtllmGen + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + config = _runtime_config() + pretrained_config = PretrainedConfig() + pretrained_config.num_experts = config.num_experts + pretrained_config.hidden_size = config.routed_expert_hidden_size + pretrained_config.intermediate_size = config.moe_intermediate_size + pretrained_config.torch_dtype = torch.bfloat16 + model_config = ModelConfig( + pretrained_config=pretrained_config, + mapping=Mapping( + world_size=2, + rank=0, + tp_size=2, + moe_tp_size=2, + moe_ep_size=1, + ), + moe_backend="TRTLLM", + ) + gate = KimiK3MoEGate(config) + moe = create_moe( + routing_method=gate.routing_method, + num_experts=config.num_experts, + hidden_size=config.routed_expert_hidden_size, + intermediate_size=config.moe_intermediate_size, + dtype=torch.bfloat16, + reduce_results=False, + model_config=model_config, + override_quant_config=QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8), + layer_idx=1, + trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, + trtllm_gen_activation_alpha=config.activation_situ_beta, + trtllm_gen_activation_beta=config.activation_situ_linear_beta, + communication_method=None, + ) + + assert isinstance(moe, ConfigurableMoE) + assert not moe.reduce_results + assert moe.all_reduce is not None + + @pytest.mark.parametrize( "attention_dp,tp_size,has_routed_comm,expected_shared_tp,expected_shared_reduce,expected_combined_ar", [ @@ -221,7 +272,7 @@ def __init__(self, *args, **kwargs): fake_moe.backend = SimpleNamespace(initial_local_expert_ids=[0, 1, 2, 3]) fake_moe.comm = object() if has_routed_comm else None fake_moe.layer_load_balancer = None - fake_moe.all_reduce = _FakeAllReduce() if not attention_dp and tp_size > 1 else None + fake_moe.all_reduce = _FakeAllReduce() monkeypatch.setattr(modeling_kimi_linear, "create_moe", lambda **_: fake_moe) monkeypatch.setattr(modeling_kimi_linear, "AllReduce", _FakeAllReduce) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py index f5bb822a6f0c..7dee839eb329 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py @@ -67,6 +67,16 @@ def test_situ_and_mul_strided_rows(): torch.testing.assert_close(out, ref, rtol=1.6e-2, atol=1e-5) +@requires_cuda +def test_situ_and_mul_rejects_strided_columns(): + """The kernel does not accept a non-contiguous packed dimension.""" + full = torch.randn(8, 1024, dtype=torch.bfloat16, device="cuda") + x = full[:, ::2] + + with pytest.raises(ValueError, match="contiguous last dimension"): + torch.ops.trtllm.situ_and_mul(x, 4.0, 25.0) + + @requires_cuda @pytest.mark.parametrize("situ_beta,situ_linear_beta", _BETAS, ids=_BETA_IDS) def test_kimi_k3_mlp_fused_activation_matches_eager(situ_beta, situ_linear_beta): From 18310323e6dcf5a7818523660275eca91e162f9b Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 13 Aug 2026 21:24:12 +0800 Subject: [PATCH 3/9] test(torch): avoid TP workspace allocation Signed-off-by: Jiagan Cheng --- tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index 57569bdec7c1..3a000d998b6b 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -190,6 +190,7 @@ def test_create_moe_keeps_all_reduce_for_external_tp_reduction(): """``reduce_results=False`` keeps the wrapper's direct-TP collective.""" from transformers.configuration_utils import PretrainedConfig + from tensorrt_llm._torch.distributed import AllReduceStrategy from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE, create_moe @@ -213,6 +214,7 @@ def test_create_moe_keeps_all_reduce_for_external_tp_reduction(): moe_ep_size=1, ), moe_backend="TRTLLM", + allreduce_strategy=AllReduceStrategy.NCCL, ) gate = KimiK3MoEGate(config) moe = create_moe( From 6a327479a5dd922d316e9e7be1eaf831e8868ebf Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 13 Aug 2026 21:33:08 +0800 Subject: [PATCH 4/9] test(torch): validate Kimi K3 MoE output Signed-off-by: Jiagan Cheng --- .../_torch/modules/moe/test_kimi_k3_mlp.py | 133 +++++++----------- 1 file changed, 52 insertions(+), 81 deletions(-) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index 3a000d998b6b..d4b164c0128f 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -185,80 +185,23 @@ def test_gated_mlp_supports_fused_situ(situ_beta, situ_linear_beta): torch.testing.assert_close(gated(x), reference(x), rtol=1.6e-2, atol=1e-3) -@requires_cuda -def test_create_moe_keeps_all_reduce_for_external_tp_reduction(): - """``reduce_results=False`` keeps the wrapper's direct-TP collective.""" - from transformers.configuration_utils import PretrainedConfig - - from tensorrt_llm._torch.distributed import AllReduceStrategy - from tensorrt_llm._torch.model_config import ModelConfig - from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate - from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE, create_moe - from tensorrt_llm._torch.utils import ActType_TrtllmGen - from tensorrt_llm.mapping import Mapping - from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig - - config = _runtime_config() - pretrained_config = PretrainedConfig() - pretrained_config.num_experts = config.num_experts - pretrained_config.hidden_size = config.routed_expert_hidden_size - pretrained_config.intermediate_size = config.moe_intermediate_size - pretrained_config.torch_dtype = torch.bfloat16 - model_config = ModelConfig( - pretrained_config=pretrained_config, - mapping=Mapping( - world_size=2, - rank=0, - tp_size=2, - moe_tp_size=2, - moe_ep_size=1, - ), - moe_backend="TRTLLM", - allreduce_strategy=AllReduceStrategy.NCCL, - ) - gate = KimiK3MoEGate(config) - moe = create_moe( - routing_method=gate.routing_method, - num_experts=config.num_experts, - hidden_size=config.routed_expert_hidden_size, - intermediate_size=config.moe_intermediate_size, - dtype=torch.bfloat16, - reduce_results=False, - model_config=model_config, - override_quant_config=QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8), - layer_idx=1, - trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, - trtllm_gen_activation_alpha=config.activation_situ_beta, - trtllm_gen_activation_beta=config.activation_situ_linear_beta, - communication_method=None, - ) - - assert isinstance(moe, ConfigurableMoE) - assert not moe.reduce_results - assert moe.all_reduce is not None - - @pytest.mark.parametrize( - "attention_dp,tp_size,has_routed_comm,expected_shared_tp,expected_shared_reduce,expected_combined_ar", + "attention_dp,tp_size,has_routed_comm", [ - (True, 2, True, 1, False, False), - (False, 1, False, 1, False, False), - (False, 2, False, 2, False, True), - (False, 2, True, 2, True, False), + (True, 2, True), + (False, 1, False), + (False, 2, False), + (False, 2, True), ], ids=["attention_dp", "single_rank", "combined_tp", "routed_comm"], ) -def test_kimi_k3_shared_expert_reduction_mode( +def test_kimi_k3_moe_output_matches_reference_across_parallel_modes( monkeypatch, attention_dp, tp_size, has_routed_comm, - expected_shared_tp, - expected_shared_reduce, - expected_combined_ar, ): - """Each parallel mode selects exactly the required shared reduction.""" - from tensorrt_llm._torch import distributed + """Every parallel mode produces the same complete block output.""" from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_kimi_linear from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE @@ -268,17 +211,39 @@ def test_kimi_k3_shared_expert_reduction_mode( class _FakeAllReduce(nn.Module): def __init__(self, *args, **kwargs): super().__init__() + self.peer_contribution = None + + def forward(self, value): + return value + self.peer_contribution + + class _FakeMoE(ConfigurableMoE): + def __init__(self): + nn.Module.__init__(self) + self.backend = SimpleNamespace(initial_local_expert_ids=[0, 1, 2, 3]) + self.comm = object() if has_routed_comm else None + self.layer_load_balancer = None + self.all_reduce = _FakeAllReduce() + self.returns_complete_output = has_routed_comm or attention_dp or tp_size == 1 + + def forward(self, hidden_states, router_logits, all_rank_num_tokens=None): + scale = 11 if self.returns_complete_output else 7 + return hidden_states * scale + + class _FakeSharedExperts(nn.Module): + def __init__(self, *, config, overridden_tp_size, reduce_output, **kwargs): + super().__init__() + self.is_sharded = config.mapping.tp_size > 1 and overridden_tp_size != 1 + self.reduce_output = reduce_output - fake_moe = ConfigurableMoE.__new__(ConfigurableMoE) - torch.nn.Module.__init__(fake_moe) - fake_moe.backend = SimpleNamespace(initial_local_expert_ids=[0, 1, 2, 3]) - fake_moe.comm = object() if has_routed_comm else None - fake_moe.layer_load_balancer = None - fake_moe.all_reduce = _FakeAllReduce() + def forward(self, hidden_states): + if not self.is_sharded or self.reduce_output: + return hidden_states * 5 + return hidden_states * 2 + + fake_moe = _FakeMoE() monkeypatch.setattr(modeling_kimi_linear, "create_moe", lambda **_: fake_moe) - monkeypatch.setattr(modeling_kimi_linear, "AllReduce", _FakeAllReduce) - monkeypatch.setattr(distributed, "AllReduce", _FakeAllReduce) + monkeypatch.setattr(modeling_kimi_linear, "GatedMLP", _FakeSharedExperts) monkeypatch.setattr(torch.cuda, "Event", lambda: object()) mapping = Mapping( @@ -289,17 +254,23 @@ def __init__(self, *args, **kwargs): ) model_config = ModelConfig(mapping=mapping, quant_config=QuantConfig()) config = _runtime_config() + config.hidden_size = config.routed_expert_hidden_size = 4 runtime = modeling_kimi_linear.KimiK3MoERuntime(model_config, config, layer_idx=1) + runtime.routed_expert_down_proj = nn.Identity() + runtime.routed_expert_norm = nn.Identity() + runtime.routed_expert_up_proj = nn.Identity() + with torch.no_grad(): + runtime.gate.weight.zero_() + runtime.gate.e_score_correction_bias.zero_() + + hidden_states = torch.arange(1, 9, dtype=torch.bfloat16).reshape(2, 4) + fake_moe.all_reduce.peer_contribution = torch.cat( + (hidden_states * 3, hidden_states * 4), dim=-1 + ) - shared = runtime.shared_experts - assert isinstance(shared, GatedMLP) - assert shared.gate_up_proj.tp_size == expected_shared_tp - assert shared.down_proj.tp_size == expected_shared_tp - assert shared.down_proj.reduce_output is expected_shared_reduce - assert runtime._use_combined_all_reduce is expected_combined_ar - local_intermediate = 512 // expected_shared_tp - assert shared.gate_up_proj.weight.shape == (2 * local_intermediate, config.hidden_size) - assert shared.down_proj.weight.shape == (config.hidden_size, local_intermediate) + actual = runtime(hidden_states) + expected = hidden_states * 16 + torch.testing.assert_close(actual, expected) @pytest.mark.parametrize( From 864f74ee5d52fa0a84d82e1c76c675dcf371cc04 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 13 Aug 2026 21:34:56 +0800 Subject: [PATCH 5/9] test(torch): use portable Kimi router path Signed-off-by: Jiagan Cheng --- tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index d4b164c0128f..e49668c4ab12 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -242,6 +242,7 @@ def forward(self, hidden_states): fake_moe = _FakeMoE() + monkeypatch.setenv("KIMI_K3_ROUTER_BF16", "0") monkeypatch.setattr(modeling_kimi_linear, "create_moe", lambda **_: fake_moe) monkeypatch.setattr(modeling_kimi_linear, "GatedMLP", _FakeSharedExperts) monkeypatch.setattr(torch.cuda, "Event", lambda: object()) From 12ad4476e8e4191ec317a167cf48d1c6a5d92b5c Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Thu, 13 Aug 2026 22:40:53 +0800 Subject: [PATCH 6/9] test(torch): validate Kimi shared MLP construction Signed-off-by: Jiagan Cheng --- .../_torch/modules/moe/test_kimi_k3_mlp.py | 72 +++++++------------ 1 file changed, 27 insertions(+), 45 deletions(-) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index e49668c4ab12..2e59e47e335a 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -14,7 +14,7 @@ ``None`` (the two activation code paths); * the row-concat convention is required: swapping the halves breaks the numerics (mutation control). -* direct MoE-TP selects combined reduction for the shared and routed partials. +* shared-expert sharding and reduction follow the selected parallel mode. """ from types import SimpleNamespace @@ -186,22 +186,27 @@ def test_gated_mlp_supports_fused_situ(situ_beta, situ_linear_beta): @pytest.mark.parametrize( - "attention_dp,tp_size,has_routed_comm", + "attention_dp,tp_size,rank,has_routed_comm,expected_shared_tp,expected_shared_rank,expected_reduce", [ - (True, 2, True), - (False, 1, False), - (False, 2, False), - (False, 2, True), + (True, 8, 7, True, 1, 0, False), + (False, 1, 0, False, 1, 0, False), + (False, 8, 7, False, 8, 7, False), + (False, 8, 7, True, 8, 7, True), ], - ids=["attention_dp", "single_rank", "combined_tp", "routed_comm"], + ids=["attention_dp", "single_rank", "direct_tp", "routed_comm"], ) -def test_kimi_k3_moe_output_matches_reference_across_parallel_modes( +def test_kimi_k3_shared_expert_parallel_construction( monkeypatch, attention_dp, tp_size, + rank, has_routed_comm, + expected_shared_tp, + expected_shared_rank, + expected_reduce, ): - """Every parallel mode produces the same complete block output.""" + """Shared experts are replicated or sharded for the selected parallel mode.""" + from tensorrt_llm._torch import distributed from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_kimi_linear from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE @@ -211,10 +216,6 @@ def test_kimi_k3_moe_output_matches_reference_across_parallel_modes( class _FakeAllReduce(nn.Module): def __init__(self, *args, **kwargs): super().__init__() - self.peer_contribution = None - - def forward(self, value): - return value + self.peer_contribution class _FakeMoE(ConfigurableMoE): def __init__(self): @@ -223,55 +224,36 @@ def __init__(self): self.comm = object() if has_routed_comm else None self.layer_load_balancer = None self.all_reduce = _FakeAllReduce() - self.returns_complete_output = has_routed_comm or attention_dp or tp_size == 1 - - def forward(self, hidden_states, router_logits, all_rank_num_tokens=None): - scale = 11 if self.returns_complete_output else 7 - return hidden_states * scale - - class _FakeSharedExperts(nn.Module): - def __init__(self, *, config, overridden_tp_size, reduce_output, **kwargs): - super().__init__() - self.is_sharded = config.mapping.tp_size > 1 and overridden_tp_size != 1 - self.reduce_output = reduce_output - - def forward(self, hidden_states): - if not self.is_sharded or self.reduce_output: - return hidden_states * 5 - return hidden_states * 2 fake_moe = _FakeMoE() monkeypatch.setenv("KIMI_K3_ROUTER_BF16", "0") monkeypatch.setattr(modeling_kimi_linear, "create_moe", lambda **_: fake_moe) - monkeypatch.setattr(modeling_kimi_linear, "GatedMLP", _FakeSharedExperts) + monkeypatch.setattr(distributed, "AllReduce", _FakeAllReduce) monkeypatch.setattr(torch.cuda, "Event", lambda: object()) mapping = Mapping( world_size=tp_size, - rank=0, + rank=rank, tp_size=tp_size, enable_attention_dp=attention_dp, ) model_config = ModelConfig(mapping=mapping, quant_config=QuantConfig()) config = _runtime_config() - config.hidden_size = config.routed_expert_hidden_size = 4 runtime = modeling_kimi_linear.KimiK3MoERuntime(model_config, config, layer_idx=1) - runtime.routed_expert_down_proj = nn.Identity() - runtime.routed_expert_norm = nn.Identity() - runtime.routed_expert_up_proj = nn.Identity() - with torch.no_grad(): - runtime.gate.weight.zero_() - runtime.gate.e_score_correction_bias.zero_() - hidden_states = torch.arange(1, 9, dtype=torch.bfloat16).reshape(2, 4) - fake_moe.all_reduce.peer_contribution = torch.cat( - (hidden_states * 3, hidden_states * 4), dim=-1 + shared = runtime.shared_experts + assert isinstance(shared, GatedMLP) + assert shared.gate_up_proj.tp_size == expected_shared_tp + assert shared.gate_up_proj.tp_rank == expected_shared_rank + assert shared.down_proj.tp_size == expected_shared_tp + assert shared.down_proj.tp_rank == expected_shared_rank + assert shared.down_proj.reduce_output is expected_reduce + local_intermediate = ( + config.moe_intermediate_size * config.num_shared_experts // expected_shared_tp ) - - actual = runtime(hidden_states) - expected = hidden_states * 16 - torch.testing.assert_close(actual, expected) + assert shared.gate_up_proj.weight.shape == (2 * local_intermediate, config.hidden_size) + assert shared.down_proj.weight.shape == (config.hidden_size, local_intermediate) @pytest.mark.parametrize( From 58bbab4fa207c238f6501b46ec8e883fa7b794f2 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Fri, 14 Aug 2026 16:07:42 +0800 Subject: [PATCH 7/9] refactor(torch): remove unused Kimi RMSNorm flag Signed-off-by: Jiagan Cheng --- .../_torch/models/modeling_kimi_linear.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 0071239d0393..895a4359de49 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -98,7 +98,6 @@ from ...models.modeling_utils import QuantAlgo, QuantConfig from ..attention_backend import AttentionMetadata from ..distributed import AllReduce, AllReduceParams, AllReduceStrategy -from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE from ..model_config import ModelConfig from ..modules.fused_moe import ConfigurableMoE, create_moe from ..modules.fused_moe.routing import DeepSeekV3MoeRoutingMethod @@ -133,10 +132,6 @@ _K3_MOE_TP_ENV = "TLLM_K3_MOE_TP_SIZE" _K3_MOE_EP_ENV = "TLLM_K3_MOE_EP_SIZE" -# Route the exact Kimi RMSNorm through FlashInfer's single-kernel path. Set -# this to "0" to retain the eager reference's byte-exact rounding behavior. -_FUSED_RMSNORM = os.environ.get("KIMI_K3_FUSED_RMSNORM", "1") == "1" - if TYPE_CHECKING: from transformers import PretrainedConfig @@ -287,16 +282,6 @@ def __init__( self.eps = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if ( - _FUSED_RMSNORM - and IS_FLASHINFER_AVAILABLE - and hidden_states.is_cuda - and hidden_states.dtype in (torch.float16, torch.bfloat16) - and self.weight.dtype == hidden_states.dtype - ): - from ..custom_ops import flashinfer_rmsnorm - - return flashinfer_rmsnorm(hidden_states.contiguous(), self.weight, self.eps) input_dtype = hidden_states.dtype hidden_states_float = hidden_states.to(torch.float32) variance = hidden_states_float.pow(2).mean(-1, keepdim=True) From 1f49d44537819b9b789ff78d2cdae4a8787285a6 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Mon, 17 Aug 2026 18:37:44 +0800 Subject: [PATCH 8/9] fix(torch): address Kimi K3 review feedback Signed-off-by: Jiagan Cheng --- .../_torch/models/modeling_kimi_linear.py | 11 +++--- .../_torch/modules/moe/test_kimi_k3_mlp.py | 17 ++++----- .../modules/moe/test_kimi_k3_moe_gate.py | 12 ------- .../modules/moe/test_kimi_k3_situ_and_mul.py | 34 ------------------ .../modules/moe/test_kimi_k3_situ_moe.py | 36 ------------------- 5 files changed, 11 insertions(+), 99 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 895a4359de49..bfbf70a34557 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -894,7 +894,6 @@ def __init__( self.experts_per_rank = len(local_expert_ids) self.expert_lo = local_expert_ids[0] self.expert_hi = self.expert_lo + self.experts_per_rank - routed_comm = self.routed_experts.comm shared_intermediate = cfg.moe_intermediate_size * cfg.num_shared_experts attention_dp = model_config.mapping.enable_attention_dp @@ -903,17 +902,15 @@ def __init__( # Under attention DP each rank owns different tokens, so the shared # expert is replicated (TP size 1) and must not reduce across ranks. # Direct MoE-TP leaves both branches as partials for one concatenated - # all-reduce. Communication-backed MoE already returns a complete - # routed result, so GatedMLP must reduce only its shared partial. + # all-reduce. use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1 - needs_routed_all_reduce = use_shared_tp and routed_comm is None routed_all_reduce = self.routed_experts.all_reduce - if needs_routed_all_reduce and routed_all_reduce is None: + if use_shared_tp and routed_all_reduce is None: raise RuntimeError( "Kimi K3 direct MoE tensor parallelism requires the " "ConfigurableMoE all-reduce even when reduce_results=False." ) - self._use_combined_all_reduce = needs_routed_all_reduce + self._use_combined_all_reduce = use_shared_tp self.shared_experts = GatedMLP( hidden_size=cfg.hidden_size, intermediate_size=shared_intermediate, @@ -926,7 +923,7 @@ def __init__( dtype=dtype, config=shared_model_config, overridden_tp_size=1 if attention_dp else None, - reduce_output=use_shared_tp and not self._use_combined_all_reduce, + reduce_output=False, layer_idx=layer_idx, is_shared_expert=True, ) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index 2e59e47e335a..6971ca85a601 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -186,24 +186,21 @@ def test_gated_mlp_supports_fused_situ(situ_beta, situ_linear_beta): @pytest.mark.parametrize( - "attention_dp,tp_size,rank,has_routed_comm,expected_shared_tp,expected_shared_rank,expected_reduce", + "attention_dp,tp_size,rank,expected_shared_tp,expected_shared_rank", [ - (True, 8, 7, True, 1, 0, False), - (False, 1, 0, False, 1, 0, False), - (False, 8, 7, False, 8, 7, False), - (False, 8, 7, True, 8, 7, True), + (True, 8, 7, 1, 0), + (False, 1, 0, 1, 0), + (False, 8, 7, 8, 7), ], - ids=["attention_dp", "single_rank", "direct_tp", "routed_comm"], + ids=["attention_dp", "single_rank", "direct_tp"], ) def test_kimi_k3_shared_expert_parallel_construction( monkeypatch, attention_dp, tp_size, rank, - has_routed_comm, expected_shared_tp, expected_shared_rank, - expected_reduce, ): """Shared experts are replicated or sharded for the selected parallel mode.""" from tensorrt_llm._torch import distributed @@ -221,7 +218,7 @@ class _FakeMoE(ConfigurableMoE): def __init__(self): nn.Module.__init__(self) self.backend = SimpleNamespace(initial_local_expert_ids=[0, 1, 2, 3]) - self.comm = object() if has_routed_comm else None + self.comm = None self.layer_load_balancer = None self.all_reduce = _FakeAllReduce() @@ -248,7 +245,7 @@ def __init__(self): assert shared.gate_up_proj.tp_rank == expected_shared_rank assert shared.down_proj.tp_size == expected_shared_tp assert shared.down_proj.tp_rank == expected_shared_rank - assert shared.down_proj.reduce_output is expected_reduce + assert shared.down_proj.reduce_output is False local_intermediate = ( config.moe_intermediate_size * config.num_shared_experts // expected_shared_tp ) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py index a887c98723d8..279b9ffcd19f 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py @@ -71,15 +71,3 @@ def test_fused_routing_matches_eager_reference(): rtol=2e-3, atol=2e-3, ) - - -@pytest.mark.parametrize( - "cfg,error", - [ - (_GateConfig(moe_router_activation_func="softmax"), "requires sigmoid"), - (_GateConfig(moe_renormalize=False), "requires top-k weight renormalization"), - ], -) -def test_invalid_config_rejected_by_production_routing(cfg, error): - with pytest.raises(ValueError, match=error): - _ = KimiK3MoEGate(cfg).routing_method diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py index 7dee839eb329..8366e87677c8 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py @@ -8,8 +8,6 @@ * elementwise parity across shapes (masked tails, multi-block rows), ``beta`` / ``linear_beta`` settings (incl. ``linear_beta=None``), bf16 in/out, tight tolerance; -* the ``KimiK3MLP(use_fused_activation=True)`` wiring matches the eager - module bit-for-bit at the down_proj output tolerance; * the fake/meta registration produces the right shape/dtype (graph tracing contract); * the op is CUDA-graph-capturable (no host sync, no data-dependent @@ -77,38 +75,6 @@ def test_situ_and_mul_rejects_strided_columns(): torch.ops.trtllm.situ_and_mul(x, 4.0, 25.0) -@requires_cuda -@pytest.mark.parametrize("situ_beta,situ_linear_beta", _BETAS, ids=_BETA_IDS) -def test_kimi_k3_mlp_fused_activation_matches_eager(situ_beta, situ_linear_beta): - device = torch.device("cuda") - hidden_size, intermediate_size = 512, 384 - torch.manual_seed(3) - eager = KimiK3MLP( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - dtype=torch.bfloat16, - device=device, - ) - with torch.no_grad(): - for proj in (eager.gate_up_proj, eager.down_proj): - proj.weight.copy_(torch.randn_like(proj.weight, dtype=torch.float32) * 0.05) - fused = KimiK3MLP( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - situ_beta=situ_beta, - situ_linear_beta=situ_linear_beta, - use_fused_activation=True, - dtype=torch.bfloat16, - device=device, - ) - fused.load_state_dict(eager.state_dict()) - - x = torch.randn(64, hidden_size, dtype=torch.bfloat16, device=device) * 0.5 - torch.testing.assert_close(fused(x), eager(x), rtol=1.6e-2, atol=1e-3) - - def test_kimi_k3_mlp_rejects_fused_flag_with_custom_activation(): with pytest.raises(ValueError, match="use_fused_activation"): KimiK3MLP( 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 1e0648ccca88..11ef6d557e50 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 @@ -170,42 +170,6 @@ def test_padded_fused_shapes(): assert padded_fused_shapes(2880, 96) == (3072, 2944, 128) -def test_kimi_gate_reuses_deepseek_v3_routing(): - config = _K3Config(num_experts=16, num_experts_per_token=4) - gate = KimiK3MoEGate(config) - torch.manual_seed(23) - with torch.no_grad(): - gate.weight.normal_(std=0.1) - gate.e_score_correction_bias.normal_(std=0.05) - hidden_states = torch.randn(2, 7, config.hidden_size) - - logits = gate.compute_logits(hidden_states) - scores = logits.sigmoid() - scores_for_choice = scores + gate.e_score_correction_bias.unsqueeze(0) - expected_ids = scores_for_choice.topk(gate.top_k, dim=-1, sorted=False).indices - expected_weights = scores.gather(1, expected_ids) - expected_weights /= expected_weights.sum(dim=-1, keepdim=True) - expected_weights *= gate.routed_scaling_factor - - routing_method = gate.routing_method - # Exercise the portable PyTorch short path; the production path keeps - # is_fused=True and uses the same routing contract. - routing_method.routing_impl.is_fused = False - actual_ids, actual_weights = routing_method.apply(logits) - - expected_order = expected_ids.argsort(dim=-1) - actual_order = actual_ids.argsort(dim=-1) - assert actual_ids.dtype == torch.int32 - torch.testing.assert_close( - expected_ids.gather(1, expected_order).to(actual_ids.dtype), - actual_ids.gather(1, actual_order), - ) - torch.testing.assert_close( - expected_weights.gather(1, expected_order), - actual_weights.gather(1, actual_order), - ) - - def test_communication_factory_accepts_model_selected_method(monkeypatch): mapping = SimpleNamespace( enable_attention_dp=True, From 63116759579862ccc9e145417a0bdfc346fa2657 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Mon, 17 Aug 2026 19:21:14 +0800 Subject: [PATCH 9/9] test(torch): select Kimi MoE backend explicitly Signed-off-by: Jiagan Cheng --- tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index 6971ca85a601..9a7e053082af 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -235,7 +235,11 @@ def __init__(self): tp_size=tp_size, enable_attention_dp=attention_dp, ) - model_config = ModelConfig(mapping=mapping, quant_config=QuantConfig()) + model_config = ModelConfig( + mapping=mapping, + quant_config=QuantConfig(), + moe_backend="TRTLLM", + ) config = _runtime_config() runtime = modeling_kimi_linear.KimiK3MoERuntime(model_config, config, layer_idx=1)