diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 0b102e49e790..bfbf70a34557 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,15 @@
 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 ..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
@@ -188,6 +197,98 @@
 _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:
+        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 +403,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 +429,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 +454,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 +501,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 +568,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 +579,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 +839,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,
@@ -768,27 +895,41 @@ def __init__(
         self.expert_lo = local_expert_ids[0]
         self.expert_hi = self.expert_lo + self.experts_per_rank
 
-        # 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.
+        use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1
+        routed_all_reduce = self.routed_experts.all_reduce
+        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 = use_shared_tp
+        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=False,
+            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 +946,14 @@ def __init__(
             hidden_size=self.moe_hidden_size, eps=cfg.rms_norm_eps, dtype=dtype
         )
 
+    @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, projection.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 +1063,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 if self._use_combined_all_reduce else None
 
         def _routed_output():
             # Latent down/up projections via the min-latency fused GEMM op:
@@ -926,37 +1076,23 @@ 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,
                 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
+            # Communication-backed paths return a complete routed result.
             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_projection(y, self.routed_expert_up_proj)
 
-        # 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 +1101,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_projection(routed_latent, self.routed_expert_up_proj)
         return routed_out + shared_out
 
 
@@ -2074,37 +2221,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 +2322,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 +2341,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 +2625,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 +2658,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 +2780,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 +2789,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..980893271bfe
--- /dev/null
+++ b/tensorrt_llm/_torch/modules/situ.py
@@ -0,0 +1,116 @@
+# 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
+    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)
+
+    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/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/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..9a7e053082af 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).
+* shared-expert sharding and reduction follow the selected parallel mode.
 """
 
+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,204 @@ 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,rank,expected_shared_tp,expected_shared_rank",
+    [
+        (True, 8, 7, 1, 0),
+        (False, 1, 0, 1, 0),
+        (False, 8, 7, 8, 7),
+    ],
+    ids=["attention_dp", "single_rank", "direct_tp"],
+)
+def test_kimi_k3_shared_expert_parallel_construction(
+    monkeypatch,
+    attention_dp,
+    tp_size,
+    rank,
+    expected_shared_tp,
+    expected_shared_rank,
+):
+    """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
+    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__()
+
+    class _FakeMoE(ConfigurableMoE):
+        def __init__(self):
+            nn.Module.__init__(self)
+            self.backend = SimpleNamespace(initial_local_expert_ids=[0, 1, 2, 3])
+            self.comm = None
+            self.layer_load_balancer = None
+            self.all_reduce = _FakeAllReduce()
+
+    fake_moe = _FakeMoE()
+
+    monkeypatch.setenv("KIMI_K3_ROUTER_BF16", "0")
+    monkeypatch.setattr(modeling_kimi_linear, "create_moe", lambda **_: fake_moe)
+    monkeypatch.setattr(distributed, "AllReduce", _FakeAllReduce)
+    monkeypatch.setattr(torch.cuda, "Event", lambda: object())
+
+    mapping = Mapping(
+        world_size=tp_size,
+        rank=rank,
+        tp_size=tp_size,
+        enable_attention_dp=attention_dp,
+    )
+    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)
+
+    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 False
+    local_intermediate = (
+        config.moe_intermediate_size * config.num_shared_experts // 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..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
@@ -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
@@ -95,33 +71,3 @@ def test_fused_routing_matches_eager():
         rtol=2e-3,
         atol=2e-3,
     )
-
-
-@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",
-    [
-        _GateConfig(moe_router_activation_func="softmax"),
-        _GateConfig(num_expert_group=4, topk_group=2),
-        _GateConfig(moe_renormalize=False),
-        _GateConfig(num_experts_per_token=1),
-    ],
-)
-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
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..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
@@ -18,8 +16,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")
 
@@ -67,35 +66,13 @@ def test_situ_and_mul_strided_rows():
 
 
 @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_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)
 
 
 def test_kimi_k3_mlp_rejects_fused_flag_with_custom_activation():
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..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
@@ -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
@@ -171,35 +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)
-
-    expected_ids, expected_weights = gate(hidden_states)
-    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))
-
-    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,