From 3a4152b1fdc74fb93b8db98ab2b3f18b0328fd82 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 28 Jul 2026 03:18:22 -0700
Subject: [PATCH 01/15] [None][feat] add Kimi K3 SiTU MegaMoE support

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 .../_torch/models/modeling_kimi_linear.py     |  39 ++-
 .../fused_moe/mega_moe/mega_moe_deepgemm.py   |  64 ++++-
 .../modules/moe/test_kimi_k3_situ_moe.py      | 255 +++++++++++++-----
 .../_torch/modules/moe/test_moe_backend.py    |  35 +++
 4 files changed, 313 insertions(+), 80 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 913e1f618339..fe504e30d23f 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -710,7 +710,7 @@ def _swap(parent: nn.Module, attr: str) -> None:
 
 
 class KimiK3MoERuntime(nn.Module):
-    """Kimi K3 latent MoE block backed by ConfigurableMoE/TRTLLM-Gen."""
+    """Kimi K3 latent MoE block backed by ConfigurableMoE."""
 
     def __init__(
         self,
@@ -751,7 +751,7 @@ def __init__(
 
         routed_moe_model_config = self._routed_moe_model_config(model_config)
         routed_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8)
-        self.routed_experts = create_moe(
+        routed_moe_kwargs = dict(
             routing_method=self.gate.routing_method,
             num_experts=self.num_experts,
             hidden_size=self.moe_hidden_size,
@@ -765,20 +765,32 @@ def __init__(
             model_config=routed_moe_model_config,
             override_quant_config=routed_quant_config,
             layer_idx=layer_idx,
-            trtllm_gen_activation_type=ActType_TrtllmGen.SiTu,
-            # Cubin alpha is the gate-side SiTU beta; cubin beta is the
-            # linear-side SiTU beta.
-            trtllm_gen_activation_alpha=float(situ_beta),
-            trtllm_gen_activation_beta=float(
-                situ_linear_beta if situ_linear_beta is not None else 1.0
-            ),
             # Let CommunicationFactory select the best available strategy.
             communication_method=None,
         )
+        if routed_moe_model_config.moe_backend == "TRTLLM":
+            routed_moe_kwargs.update(
+                trtllm_gen_activation_type=ActType_TrtllmGen.SiTu,
+                # Cubin alpha is the gate-side SiTU beta; cubin beta is the
+                # linear-side SiTU beta.
+                trtllm_gen_activation_alpha=float(situ_beta),
+                trtllm_gen_activation_beta=float(
+                    situ_linear_beta if situ_linear_beta is not None else 1.0
+                ),
+            )
+        self.routed_experts = create_moe(**routed_moe_kwargs)
         if not isinstance(self.routed_experts, ConfigurableMoE):
             raise RuntimeError(
                 "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled."
             )
+        if routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM":
+            from ..modules.fused_moe.mega_moe import MegaMoEDeepGemm
+
+            if not isinstance(self.routed_experts.backend, MegaMoEDeepGemm):
+                raise RuntimeError(
+                    "Kimi K3 explicitly requested MEGAMOE_DEEPGEMM, but the "
+                    f"MoE factory selected {type(self.routed_experts.backend).__name__}."
+                )
         if self.routed_experts.layer_load_balancer is not None:
             raise NotImplementedError(
                 "Kimi K3 packed-checkpoint streaming does not yet support "
@@ -801,6 +813,7 @@ def __init__(
             not _K3_DISABLE_MIN_LATENCY_LATENT_PROJ
             and not _K3_DISABLE_FUSED_LATENT_DOWN_MXFP8
             and routed_comm is None
+            and hasattr(self.routed_experts.backend, "op_backend")
         )
 
         # Shared experts stay replicated (DeepSeek's attention-DP
@@ -915,7 +928,13 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig:
         routed_model_config._frozen = False
         routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs)
         routed_model_config.mapping = routed_mapping
-        routed_model_config.moe_backend = "TRTLLM"
+        # Preserve K3's TRTLLM-Gen default while allowing the explicit
+        # DeepGEMM MegaMoE opt-in used by the SM100 FP8xFP4 SiTU path.
+        routed_model_config.moe_backend = (
+            "MEGAMOE_DEEPGEMM"
+            if model_config.moe_backend.upper() == "MEGAMOE_DEEPGEMM"
+            else "TRTLLM"
+        )
         routed_model_config._frozen = True
         return routed_model_config
 
diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
index 23b82228d196..132d4ccd3f26 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
@@ -118,7 +118,7 @@ class MegaMoEDeepGemm(MoE):
 
     _SUPPORTED_ACTIVATION_DTYPES = frozenset({torch.bfloat16})
 
-    # Kernel owns dispatch + GEMM1 + SwiGLU + GEMM2 + combine via NVLink
+    # Kernel owns dispatch + GEMM1 + gated activation + GEMM2 + combine via NVLink
     # SymmBuffer; ConfigurableMoE must NOT layer host-side comm on top.
     scheduler_kind = MoESchedulerKind.FUSED_COMM
 
@@ -203,11 +203,13 @@ def __init__(
         layer_idx: Optional[int] = None,
         activation_type: ActivationType = ActivationType.Swiglu,
         init_load_balancer: bool = True,
-        # DG tunables. ``swiglu_limit_scalar`` mirrors the upstream MoE
-        # kwarg; bridged to DG's ``activation_clamp`` at the call site.
-        activation: str = "swiglu",
+        # DG tunables. ``activation=None`` infers Kimi K3 SiTU from the
+        # pretrained config and otherwise defaults to SwiGLU.
+        activation: Optional[str] = None,
         swiglu_limit_scalar: Optional[float] = None,
         fast_math: bool = True,
+        situ_beta: Optional[float] = None,
+        situ_linear_beta: Optional[float] = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -276,18 +278,27 @@ def __init__(
             "not equivalent. Use a different MoE backend for models that "
             "require pre-scaling, or extend the kernel call."
         )
-        # DG's fp8_fp4_mega_moe currently only ships a fused SwiGLU
-        # activation path. Reject other ActivationType values explicitly so
-        # ``create_moe_backend`` callers do not silently get SwiGLU when
-        # they asked for GELU / etc.
+        # ``ActivationType.Swiglu`` describes the gated FC1 tensor geometry
+        # shared by SwiGLU and SiTU. The DeepGEMM-specific activation selects
+        # the actual elementwise function below.
         if activation_type != ActivationType.Swiglu:
             raise ValueError(
                 f"MegaMoEDeepGemm only supports ActivationType.Swiglu (got {activation_type})."
             )
+        activation, situ_beta, situ_linear_beta = self._resolve_activation_config(
+            model_config,
+            activation=activation,
+            situ_beta=situ_beta,
+            situ_linear_beta=situ_linear_beta,
+        )
+        if activation == "situ" and swiglu_limit_scalar is not None:
+            raise ValueError("MegaMoEDeepGemm SiTU does not support activation_clamp.")
         self.apply_router_weight_on_input = apply_router_weight_on_input
         self.activation = activation
         self.swiglu_limit_scalar = swiglu_limit_scalar
         self.fast_math = fast_math
+        self.situ_beta = situ_beta
+        self.situ_linear_beta = situ_linear_beta
 
         # Buffer sizing. MoE layers execute serially per forward; a single
         # process-level pool sized to worst-case per-rank tokens serves all.
@@ -347,6 +358,41 @@ def __init__(
         if not model_config.skip_create_weights_in_init:
             self.create_weights()
 
+    @staticmethod
+    def _resolve_activation_config(
+        model_config: ModelConfig,
+        *,
+        activation: Optional[str],
+        situ_beta: Optional[float],
+        situ_linear_beta: Optional[float],
+    ) -> Tuple[str, Optional[float], Optional[float]]:
+        pretrained_config = model_config.pretrained_config
+        config_situ_beta = getattr(pretrained_config, "activation_situ_beta", None)
+        config_situ_linear_beta = getattr(pretrained_config, "activation_situ_linear_beta", None)
+        if activation is None:
+            activation = "situ" if config_situ_beta is not None else "swiglu"
+        activation = activation.lower()
+        if activation not in ("swiglu", "situ"):
+            raise ValueError(
+                f"MegaMoEDeepGemm activation must be 'swiglu' or 'situ'; got {activation!r}."
+            )
+        if activation == "swiglu":
+            if situ_beta is not None or situ_linear_beta is not None:
+                raise ValueError("SiTU beta parameters require activation='situ'.")
+            return activation, None, None
+
+        situ_beta = config_situ_beta if situ_beta is None else situ_beta
+        situ_linear_beta = config_situ_linear_beta if situ_linear_beta is None else situ_linear_beta
+        if situ_beta is None or situ_linear_beta is None:
+            raise ValueError(
+                "MegaMoEDeepGemm SiTU requires activation_situ_beta and "
+                "activation_situ_linear_beta in the pretrained config, or "
+                "explicit situ_beta and situ_linear_beta arguments."
+            )
+        if situ_beta <= 0 or situ_linear_beta <= 0:
+            raise ValueError("MegaMoEDeepGemm SiTU beta parameters must be positive.")
+        return activation, float(situ_beta), float(situ_linear_beta)
+
     def _supports_load_balancer(self) -> bool:
         # The DeepGEMM mega kernel routes by `topk_idx` interpreted as slot id
         # (range [0, num_slots)) once the SymmBuffer is sized to num_slots.
@@ -711,5 +757,7 @@ def run_moe(
             activation=self.activation,
             activation_clamp=self.swiglu_limit_scalar,
             fast_math=self.fast_math,
+            situ_beta=self.situ_beta,
+            situ_linear_beta=self.situ_linear_beta,
         )
         return y.to(output_dtype)
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 6c716f2bccba..874f1aad378c 100644
--- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py
+++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py
@@ -26,6 +26,7 @@
 
 import pytest
 import torch
+import torch.distributed as dist
 from utils.util import check_accuracy
 
 from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory
@@ -466,6 +467,34 @@ def test_kimi_k3_moe_split_selection(monkeypatch):
     assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2)
 
 
+def test_kimi_k3_routed_config_preserves_explicit_megamoe_backend():
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoERuntime
+    from tensorrt_llm.mapping import Mapping
+
+    model_config = ModelConfig(
+        mapping=Mapping(world_size=1, rank=0, tp_size=1),
+        moe_backend="MEGAMOE_DEEPGEMM",
+    )
+
+    routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config)
+
+    assert routed_model_config.moe_backend == "MEGAMOE_DEEPGEMM"
+    assert model_config.moe_backend == "MEGAMOE_DEEPGEMM"
+
+
+def test_kimi_k3_routed_config_keeps_trtllm_default():
+    from tensorrt_llm._torch.model_config import ModelConfig
+    from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoERuntime
+    from tensorrt_llm.mapping import Mapping
+
+    model_config = ModelConfig(mapping=Mapping(world_size=1, rank=0, tp_size=1))
+
+    routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config)
+
+    assert routed_model_config.moe_backend == "TRTLLM"
+
+
 # ---------------------------------------------------------------------------
 # MoE tensor-parallel shard parity (ConfigurableMoE / TRTLLM-Gen, GPU).
 #
@@ -496,18 +525,20 @@ def scales(*shape):
 
     bank = []
     for _ in range(num_experts):
-        bank.append({
-            "w1": nibbles(intermediate, hidden // 2),
-            "w1_sf": scales(intermediate, hidden // 32),
-            "w3": nibbles(intermediate, hidden // 2),
-            "w3_sf": scales(intermediate, hidden // 32),
-            "w2": nibbles(hidden, intermediate // 2),
-            "w2_sf": scales(hidden, intermediate // 32),
-        })
+        bank.append(
+            {
+                "w1": nibbles(intermediate, hidden // 2),
+                "w1_sf": scales(intermediate, hidden // 32),
+                "w3": nibbles(intermediate, hidden // 2),
+                "w3_sf": scales(intermediate, hidden // 32),
+                "w2": nibbles(hidden, intermediate // 2),
+                "w2_sf": scales(hidden, intermediate // 32),
+            }
+        )
     return bank
 
 
-def _make_test_gate(num_experts=_TP_EXPERTS, seed=71):
+def _make_test_gate(num_experts=_TP_EXPERTS, top_k=_TP_TOPK, seed=71):
     """One deterministically-initialized gate SHARED by all modules under
     comparison: the fused routing kernel applies the gate's
     e_score_correction_bias per module, so a per-module `torch.empty`
@@ -516,21 +547,25 @@ def _make_test_gate(num_experts=_TP_EXPERTS, seed=71):
     cfg = _K3Config(
         hidden_size=_TP_HIDDEN,
         num_experts=num_experts,
-        num_experts_per_token=min(_TP_TOPK, num_experts),
+        num_experts_per_token=min(top_k, num_experts),
     )
     gate = KimiK3MoEGate(cfg)
     gen = torch.Generator().manual_seed(seed)
     with torch.no_grad():
-        gate.weight.copy_(
-            torch.randn(gate.weight.shape, generator=gen,
-                        dtype=torch.float32) * 0.05)
+        gate.weight.copy_(torch.randn(gate.weight.shape, generator=gen, dtype=torch.float32) * 0.05)
         gate.e_score_correction_bias.copy_(
-            torch.randn(gate.e_score_correction_bias.shape, generator=gen,
-                        dtype=torch.float32) * 0.1)
+            torch.randn(gate.e_score_correction_bias.shape, generator=gen, dtype=torch.float32)
+            * 0.1
+        )
     return gate.cuda()
 
 
-def _make_routed_moe(intermediate_size, gate, num_experts=_TP_EXPERTS):
+def _make_routed_moe(
+    intermediate_size,
+    gate,
+    num_experts=_TP_EXPERTS,
+    moe_backend="TRTLLM",
+):
     """Mirror KimiK3MoERuntime's create_moe call on a single-rank mapping."""
     from transformers.configuration_utils import PretrainedConfig
 
@@ -544,12 +579,14 @@ def _make_routed_moe(intermediate_size, gate, num_experts=_TP_EXPERTS):
     pretrained_config.hidden_size = _TP_HIDDEN
     pretrained_config.intermediate_size = intermediate_size
     pretrained_config.torch_dtype = torch.bfloat16
+    pretrained_config.activation_situ_beta = 4.0
+    pretrained_config.activation_situ_linear_beta = 25.0
     model_config = ModelConfig(
         pretrained_config=pretrained_config,
         mapping=Mapping(),
-        moe_backend="TRTLLM",
+        moe_backend=moe_backend,
     )
-    moe = create_moe(
+    moe_kwargs = dict(
         routing_method=gate.routing_method,
         num_experts=num_experts,
         hidden_size=_TP_HIDDEN,
@@ -559,11 +596,15 @@ def _make_routed_moe(intermediate_size, gate, num_experts=_TP_EXPERTS):
         model_config=model_config,
         override_quant_config=QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8),
         layer_idx=0,
-        trtllm_gen_activation_type=ActType_TrtllmGen.SiTu,
-        trtllm_gen_activation_alpha=4.0,
-        trtllm_gen_activation_beta=25.0,
         communication_method=None,
-    ).cuda()
+    )
+    if moe_backend == "TRTLLM":
+        moe_kwargs.update(
+            trtllm_gen_activation_type=ActType_TrtllmGen.SiTu,
+            trtllm_gen_activation_alpha=4.0,
+            trtllm_gen_activation_beta=25.0,
+        )
+    moe = create_moe(**moe_kwargs).cuda()
     assert isinstance(moe, ConfigurableMoE)
     return moe
 
@@ -577,29 +618,45 @@ def _load_bank(moe, bank, tp_size=1, tp_rank=0):
     shard-sized parameters.
     """
     backend = moe.backend
-    proxy = SimpleNamespace(
-        expert_size_per_partition=backend.expert_size_per_partition,
-        initial_local_expert_ids=backend.initial_local_expert_ids,
-        scaling_vector_size=backend.scaling_vector_size,
-        tp_size=tp_size,
-        tp_rank=tp_rank,
-        w3_w1_weight=backend.w3_w1_weight,
-        w2_weight=backend.w2_weight,
-        w3_w1_weight_scale=backend.w3_w1_weight_scale,
-        w2_weight_scale=backend.w2_weight_scale,
-    )
-    for expert_id, tensors in enumerate(bank):
-        backend.quant_method.load_packed_mxfp4_expert(
-            proxy,
-            global_expert_id=expert_id,
-            local_slot_id=expert_id,
-            w1_weight=tensors["w1"],
-            w1_weight_scale=tensors["w1_sf"],
-            w2_weight=tensors["w2"],
-            w2_weight_scale=tensors["w2_sf"],
-            w3_weight=tensors["w3"],
-            w3_weight_scale=tensors["w3_sf"],
+    if hasattr(backend.quant_method, "load_packed_mxfp4_expert"):
+        proxy = SimpleNamespace(
+            expert_size_per_partition=backend.expert_size_per_partition,
+            initial_local_expert_ids=backend.initial_local_expert_ids,
+            scaling_vector_size=backend.scaling_vector_size,
+            tp_size=tp_size,
+            tp_rank=tp_rank,
+            w3_w1_weight=backend.w3_w1_weight,
+            w2_weight=backend.w2_weight,
+            w3_w1_weight_scale=backend.w3_w1_weight_scale,
+            w2_weight_scale=backend.w2_weight_scale,
         )
+        for expert_id, tensors in enumerate(bank):
+            backend.quant_method.load_packed_mxfp4_expert(
+                proxy,
+                global_expert_id=expert_id,
+                local_slot_id=expert_id,
+                w1_weight=tensors["w1"],
+                w1_weight_scale=tensors["w1_sf"],
+                w2_weight=tensors["w2"],
+                w2_weight_scale=tensors["w2_sf"],
+                w3_weight=tensors["w3"],
+                w3_weight_scale=tensors["w3_sf"],
+            )
+    else:
+        weights = {}
+        for expert_id, tensors in enumerate(bank):
+            prefix = f"{expert_id}."
+            weights.update(
+                {
+                    prefix + "w1.weight": tensors["w1"],
+                    prefix + "w1.weight_scale": tensors["w1_sf"],
+                    prefix + "w2.weight": tensors["w2"],
+                    prefix + "w2.weight_scale": tensors["w2_sf"],
+                    prefix + "w3.weight": tensors["w3"],
+                    prefix + "w3.weight_scale": tensors["w3_sf"],
+                }
+            )
+        backend.load_weights([weights])
     backend._weights_transformed = False
     moe.post_load_weights()
     return moe
@@ -626,23 +683,24 @@ def test_tp_shard_loader_matches_manual_slice(tp_size):
         rows = slice(tp_rank * ipp, (tp_rank + 1) * ipp)
         cols_packed = slice(tp_rank * (ipp // 2), (tp_rank + 1) * (ipp // 2))
         cols_sf = slice(tp_rank * (ipp // 32), (tp_rank + 1) * (ipp // 32))
-        manual_bank = [{
-            "w1": e["w1"][rows].contiguous(),
-            "w1_sf": e["w1_sf"][rows].contiguous(),
-            "w3": e["w3"][rows].contiguous(),
-            "w3_sf": e["w3_sf"][rows].contiguous(),
-            "w2": e["w2"][:, cols_packed].contiguous(),
-            "w2_sf": e["w2_sf"][:, cols_sf].contiguous(),
-        } for e in bank]
+        manual_bank = [
+            {
+                "w1": e["w1"][rows].contiguous(),
+                "w1_sf": e["w1_sf"][rows].contiguous(),
+                "w3": e["w3"][rows].contiguous(),
+                "w3_sf": e["w3_sf"][rows].contiguous(),
+                "w2": e["w2"][:, cols_packed].contiguous(),
+                "w2_sf": e["w2_sf"][:, cols_sf].contiguous(),
+            }
+            for e in bank
+        ]
         via_manual = _make_routed_moe(ipp, gate, num_experts=num_experts)
         _load_bank(via_manual, manual_bank)
 
-        for name in ("w3_w1_weight", "w2_weight", "w3_w1_weight_scale",
-                     "w2_weight_scale"):
+        for name in ("w3_w1_weight", "w2_weight", "w3_w1_weight_scale", "w2_weight_scale"):
             a = getattr(via_shard.backend, name).data
             b = getattr(via_manual.backend, name).data
-            assert torch.equal(a, b), (
-                f"{name} mismatch for tp_size={tp_size} tp_rank={tp_rank}")
+            assert torch.equal(a, b), f"{name} mismatch for tp_size={tp_size} tp_rank={tp_rank}"
 
 
 @situ_supported
@@ -663,14 +721,12 @@ def test_tp8_sharded_forward_matches_whole_expert(num_tokens):
     _load_bank(whole, bank)
 
     torch.manual_seed(3)
-    x = torch.randn(
-        num_tokens, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5
+    x = torch.randn(num_tokens, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5
     router_logits = gate.compute_logits(x)
 
     out_whole = whole.forward(x, router_logits, all_rank_num_tokens=None)
 
-    partial_sum = torch.zeros(num_tokens, _TP_HIDDEN, dtype=torch.float32,
-                              device="cuda")
+    partial_sum = torch.zeros(num_tokens, _TP_HIDDEN, dtype=torch.float32, device="cuda")
     for tp_rank in range(tp_size):
         shard = _make_routed_moe(ipp, gate)
         _load_bank(shard, bank, tp_size=tp_size, tp_rank=tp_rank)
@@ -679,5 +735,80 @@ def test_tp8_sharded_forward_matches_whole_expert(num_tokens):
         del shard
         torch.cuda.empty_cache()
 
-    check_accuracy(partial_sum.to(torch.bfloat16), out_whole,
-                   atol=0.08, rtol=0.08, percent=0.98)
+    check_accuracy(partial_sum.to(torch.bfloat16), out_whole, atol=0.08, rtol=0.08, percent=0.98)
+
+
+@situ_supported
+@pytest.mark.parametrize("num_tokens", [1, 16, 128], ids=lambda n: f"tokens{n}")
+@pytest.mark.parametrize(
+    "num_experts,top_k",
+    [
+        pytest.param(8, 1, id="experts8-top1"),
+        pytest.param(8, 2, id="experts8-top2"),
+        pytest.param(32, 16, id="experts32-top16"),
+    ],
+)
+def test_megamoe_deepgemm_situ_matches_trtllm_gen(num_tokens, num_experts, top_k):
+    """Compare SiTU kernels with identical packed MXFP4 weights and routing.
+
+    MegaMoE folds routing weights into the FC1 activation before its MXFP8
+    requantization, while TRTLLM-Gen combines after the expert output. The
+    quantized graphs therefore need semantic, rather than elementwise,
+    parity: high cosine similarity and bounded relative L2 error.
+    """
+    if not dist.is_initialized():
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29561")
+        os.environ.setdefault("RANK", "0")
+        os.environ.setdefault("WORLD_SIZE", "1")
+        os.environ.setdefault("LOCAL_RANK", "0")
+        torch.cuda.set_device(0)
+        dist.init_process_group(
+            backend="nccl",
+            rank=0,
+            world_size=1,
+        )
+    bank = _make_packed_expert_bank(num_experts, _TP_INTERMEDIATE, _TP_HIDDEN)
+    gate = _make_test_gate(num_experts=num_experts, top_k=top_k)
+
+    trtllm_gen = _load_bank(
+        _make_routed_moe(_TP_INTERMEDIATE, gate, num_experts=num_experts),
+        bank,
+    )
+    mega_moe = _load_bank(
+        _make_routed_moe(
+            _TP_INTERMEDIATE,
+            gate,
+            num_experts=num_experts,
+            moe_backend="MEGAMOE_DEEPGEMM",
+        ),
+        bank,
+    )
+
+    torch.manual_seed(37)
+    x = torch.randn(num_tokens, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5
+    router_logits = gate.compute_logits(x)
+
+    with torch.inference_mode():
+        trtllm_gen_output = trtllm_gen(x, router_logits)
+        mega_moe_output = mega_moe(x, router_logits)
+
+    assert torch.isfinite(mega_moe_output).all()
+    diff = (mega_moe_output.float() - trtllm_gen_output.float()).abs()
+    ref = trtllm_gen_output.float()
+    cosine = torch.nn.functional.cosine_similarity(
+        mega_moe_output.float().flatten(),
+        ref.flatten(),
+        dim=0,
+    )
+    relative_l2 = torch.linalg.vector_norm(diff) / torch.linalg.vector_norm(ref)
+    print(
+        f"M={num_tokens}, experts={num_experts}, top_k={top_k}: "
+        f"cosine={cosine.item():.8f}, "
+        f"relative_l2={relative_l2.item():.8f}, "
+        f"mean_abs={diff.mean().item():.8f}, "
+        f"p99_abs={torch.quantile(diff, 0.99).item():.8f}, "
+        f"max_abs={diff.max().item():.8f}"
+    )
+    assert cosine > 0.998
+    assert relative_l2 < 0.06
diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py
index 8223e08d67bb..361f1cb7eed7 100644
--- a/tests/unittest/_torch/modules/moe/test_moe_backend.py
+++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py
@@ -467,6 +467,41 @@ def test_megamoe_deepgemm_cache_derived_state_allocates_symm_buffer():
     quant_method.cache_derived_state.assert_called_once_with(moe)
 
 
+def test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config():
+    model_config = ModelConfig(
+        pretrained_config=SimpleNamespace(
+            activation_situ_beta=4.0,
+            activation_situ_linear_beta=25.0,
+        )
+    )
+
+    activation, situ_beta, situ_linear_beta = MegaMoEDeepGemm._resolve_activation_config(
+        model_config,
+        activation=None,
+        situ_beta=None,
+        situ_linear_beta=None,
+    )
+
+    assert activation == "situ"
+    assert situ_beta == 4.0
+    assert situ_linear_beta == 25.0
+
+
+def test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config():
+    model_config = ModelConfig(pretrained_config=SimpleNamespace())
+
+    activation, situ_beta, situ_linear_beta = MegaMoEDeepGemm._resolve_activation_config(
+        model_config,
+        activation=None,
+        situ_beta=None,
+        situ_linear_beta=None,
+    )
+
+    assert activation == "swiglu"
+    assert situ_beta is None
+    assert situ_linear_beta is None
+
+
 def test_megamoe_init_rejects_uneven_num_slots_with_value_error():
     routing_method = RenormalizeMoeRoutingMethod(top_k=1)
     model_config = ModelConfig(

From e906e6bd462eff3238f01c7496e8ee865682bcc7 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Sun, 2 Aug 2026 19:13:37 -0700
Subject: [PATCH 02/15] [None][chore] map Kimi DeepGEMM fork inputs

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 3rdparty/fetch_content.json                      | 4 ++--
 scripts/attribution/data/dependency_metadata.yml | 4 ++--
 scripts/attribution/data/files_to_dependency.yml | 6 +++---
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json
index 31820f223a18..420f56269d4f 100644
--- a/3rdparty/fetch_content.json
+++ b/3rdparty/fetch_content.json
@@ -31,8 +31,8 @@
     },
     {
       "name": "deepgemm",
-      "git_repository": "https://github.com/deepseek-ai/DeepGEMM",
-      "git_tag": "f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca",
+      "git_repository": "https://github.com/longlee0622/DeepGEMM",
+      "git_tag": "c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6",
       "git_submodules_recurse": true,
       "source_subdir": "dont-add-this-project-with-add-subdirectory"
     },
diff --git a/scripts/attribution/data/dependency_metadata.yml b/scripts/attribution/data/dependency_metadata.yml
index 2b8485cb1383..46e990cbff78 100644
--- a/scripts/attribution/data/dependency_metadata.yml
+++ b/scripts/attribution/data/dependency_metadata.yml
@@ -23,9 +23,9 @@ cutlass/v4.3.0:
 deep_ep/5be51b228a7c82dbdb213ea58e77bffd12b38af8:
   license: 452b3ef002dc6ec283fb723f0dd84997
   source: https://github.com/deepseek-ai/DeepEP/tree/5be51b228a7c82dbdb213ea58e77bffd12b38af8
-deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca:
+deepgemm/c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6:
   license: 452b3ef002dc6ec283fb723f0dd84997
-  source: https://github.com/deepseek-ai/DeepGEMM/tree/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca
+  source: https://github.com/longlee0622/DeepGEMM/tree/c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6
 dlpack/v1.0:
   license: cd9881918c97ec7b4962691660bb733e
 flashmla/1408756a88e52a25196b759eaf8db89d2b51b5a1:
diff --git a/scripts/attribution/data/files_to_dependency.yml b/scripts/attribution/data/files_to_dependency.yml
index 7a5c25002735..69505bab1ca9 100644
--- a/scripts/attribution/data/files_to_dependency.yml
+++ b/scripts/attribution/data/files_to_dependency.yml
@@ -2093,7 +2093,7 @@ deep_ep/5be51b228a7c82dbdb213ea58e77bffd12b38af8:
 - dbc9ea8cf83b20e6ce4c6f383b700f29
 - e4024308dd534f83de92752e1d7cd9a8
 - f41ae95dbbafe6107dd98bf66af018ea
-deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca:
+deepgemm/c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6:
 - 07488ed395a262f652b63d0b1c1bb3a8
 - 1101099cb0a9c0489f9e4e49a719941a
 - 111cfba37978c3bad26617cb42407970
@@ -2133,7 +2133,7 @@ deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca:
 - 95b8e1175d420f5a05bc9fb0d0c55214
 - 9bf6bbf8da71d31836279a87ffafe108
 - 9e16e23b6894db65b377f93673dc733f
-- a359f41b0ec5d67c1d9fda93dc65c5b5
+- 62fe0935dc226557c42dbacf44b2dec0
 - a3645795bcf4bad7c333094975c558b6
 - a775f6a60d47cd428f4c04289d4e8cb3
 - a9f87d66fb89c05e1ed20a9459c54f68
@@ -2157,7 +2157,7 @@ deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca:
 - ea07df16c9a083277f55ae219c85f39d
 - ef5544cabdf0490063f2b2959f62a8cf
 - ef712bc72e01afbfdbceef82c6b49174
-- f43084b87a14bcf56920e1e3b1ad2dba
+- cfcffb4170384e7699c65ed604fbd1d8
 - f5a3009221d096c818cf26fbdb4d9693
 - f5d68cc5860baa1de18f4aee4ccb0cc2
 - f7f27b18dae31db0aad429bfaaf4615e

From 63fc93fa0ccfc621552aec0974b6c98a7fb63631 Mon Sep 17 00:00:00 2001
From: Xin Guan <xguan@aws-cmh-slurm-1-vscode-01.cm.cluster>
Date: Fri, 31 Jul 2026 00:15:29 -0700
Subject: [PATCH 03/15] [None][fix] Let MegaMoE load the Kimi K3 packed MXFP4
 checkpoint

Port the packed-expert streaming adapter and MegaMoE capacity fix from xguannv/TensorRT-LLM commit d39eb40590388ac809b9e8828d1550214488a5d3.

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 examples/kimi_k3/eval_extra_llm_options.yaml  |   2 +-
 .../_torch/modules/fused_moe/quantization.py  | 104 ++++++++++++++++++
 2 files changed, 105 insertions(+), 1 deletion(-)

diff --git a/examples/kimi_k3/eval_extra_llm_options.yaml b/examples/kimi_k3/eval_extra_llm_options.yaml
index 813fe8e0dbe8..c93e2f9607b1 100644
--- a/examples/kimi_k3/eval_extra_llm_options.yaml
+++ b/examples/kimi_k3/eval_extra_llm_options.yaml
@@ -11,7 +11,7 @@ cuda_graph_config:
   enable_padding: true
   max_batch_size: 32
 moe_config:
-  max_num_tokens: 33024
+  max_num_tokens: 131072
   use_low_precision_moe_combine: true
 kv_cache_config:
   enable_block_reuse: false
diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py
index e134fd26b546..9dc728fe44a6 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py
@@ -6427,6 +6427,7 @@ def create_weights(self, module: torch.nn.Module) -> None:
         # Downstream reload/EPLB metadata path; populated lazily when parameter
         # replacement records tensors that need rebuilding before reload.
         module.rebuild_tensor_metadata = {}
+        module._packed_mxfp4_loaded_slots = set()
         self.setup_quant_scales(module)
 
     def setup_quant_scales(self, module: torch.nn.Module):
@@ -6511,6 +6512,109 @@ def _load_expert_weights_to_dst(
                 w2_scale, dst_w2_weight_scale),
                                                non_blocking=True)
 
+    def load_packed_mxfp4_expert(
+        self,
+        module: torch.nn.Module,
+        *,
+        global_expert_id: int,
+        local_slot_id: int,
+        w1_weight: torch.Tensor,
+        w1_weight_scale: torch.Tensor,
+        w2_weight: torch.Tensor,
+        w2_weight_scale: torch.Tensor,
+        w3_weight: torch.Tensor,
+        w3_weight_scale: torch.Tensor,
+    ) -> None:
+        """Load one group-32 packed MXFP4 checkpoint expert into a local slot.
+
+        Kimi K3 streams its packed checkpoint expert-by-expert so that each
+        safetensors mapping stays short-lived, and therefore calls this adapter
+        instead of ``load_weights``. The bytes written here are exactly what
+        ``_load_expert_weights_to_dst`` writes; only the source differs
+        (explicit tensors rather than a weight dict).
+
+        ``transform_weights`` is not run here. The Kimi loader clears
+        ``_weights_transformed`` once every slot is filled, so the DG-native
+        layout is rebuilt lazily from the raw bytes staged below.
+        """
+        if not 0 <= local_slot_id < module.expert_size_per_partition:
+            raise IndexError(f"local_slot_id={local_slot_id} is outside "
+                             f"[0, {module.expert_size_per_partition}).")
+        expected_expert_id = module.initial_local_expert_ids[local_slot_id]
+        if global_expert_id != expected_expert_id:
+            raise ValueError(
+                f"local slot {local_slot_id} expects global expert "
+                f"{expected_expert_id}, got {global_expert_id}.")
+
+        for name, value in (
+            ("w1_weight", w1_weight),
+            ("w1_weight_scale", w1_weight_scale),
+            ("w2_weight", w2_weight),
+            ("w2_weight_scale", w2_weight_scale),
+            ("w3_weight", w3_weight),
+            ("w3_weight_scale", w3_weight_scale),
+        ):
+            if value.dtype != torch.uint8:
+                raise TypeError(
+                    f"{name} must contain packed MXFP4 uint8 data, got "
+                    f"{value.dtype}.")
+
+        loaded_slots = module._packed_mxfp4_loaded_slots
+        if local_slot_id in loaded_slots:
+            raise ValueError(
+                f"Packed MXFP4 local slot {local_slot_id} was loaded twice.")
+        if not loaded_slots:
+            # Raw weights are about to change, so any DG-derived tensors left
+            # from an earlier load must not survive into transform_weights.
+            self._clear_transformed_weight_cache(module)
+
+        dst_w3_w1_weight = module.w3_w1_weight.data
+        dst_w3_w1_weight_scale = module.w3_w1_weight_scale.data
+        dst_w2_weight = module.w2_weight.data
+        dst_w2_weight_scale = module.w2_weight_scale.data
+
+        # DeepGEMM expects L1 in [gate | up] order before
+        # transform_weights_for_mega_moe interleaves gate/up rows, and TRT-LLM
+        # checkpoints map gate_proj -> w1 and up_proj -> w3. So despite the
+        # parameter being named w3_w1_*, the concatenation is [w1 | w3] --
+        # the same order _load_expert_weights_to_dst uses. Reversing it stays
+        # shape-compatible and fails silently in the numerics, not loudly.
+        dst_w3_w1_weight[local_slot_id].copy_(
+            torch.cat(
+                [
+                    self._to_weight_device_uint8(w1_weight, dst_w3_w1_weight),
+                    self._to_weight_device_uint8(w3_weight, dst_w3_w1_weight),
+                ],
+                dim=0,
+            ),
+            non_blocking=True,
+        )
+        dst_w3_w1_weight_scale[local_slot_id].copy_(
+            torch.cat(
+                [
+                    self._to_weight_device_uint8(w1_weight_scale,
+                                                 dst_w3_w1_weight_scale),
+                    self._to_weight_device_uint8(w3_weight_scale,
+                                                 dst_w3_w1_weight_scale),
+                ],
+                dim=0,
+            ),
+            non_blocking=True,
+        )
+        dst_w2_weight[local_slot_id].copy_(
+            self._to_weight_device_uint8(w2_weight, dst_w2_weight),
+            non_blocking=True,
+        )
+        dst_w2_weight_scale[local_slot_id].copy_(
+            self._to_weight_device_uint8(w2_weight_scale, dst_w2_weight_scale),
+            non_blocking=True,
+        )
+
+        loaded_slots.add(local_slot_id)
+        # transform_weights asserts this; the streaming path never reaches
+        # load_weights, which is where it would otherwise be set.
+        module._weights_loaded = True
+
     def load_weights(
         self,
         module: torch.nn.Module,

From 0ef1d48f6b70e103e0037167607cdc8b0a820097 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Sun, 2 Aug 2026 20:08:18 -0700
Subject: [PATCH 04/15] [None][test] support MegaMoE packed weight fixture

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 .../modules/moe/test_kimi_k3_situ_moe.py      | 26 ++++++++++---------
 1 file changed, 14 insertions(+), 12 deletions(-)

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 874f1aad378c..27c6044852eb 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
@@ -619,20 +619,22 @@ def _load_bank(moe, bank, tp_size=1, tp_rank=0):
     """
     backend = moe.backend
     if hasattr(backend.quant_method, "load_packed_mxfp4_expert"):
-        proxy = SimpleNamespace(
-            expert_size_per_partition=backend.expert_size_per_partition,
-            initial_local_expert_ids=backend.initial_local_expert_ids,
-            scaling_vector_size=backend.scaling_vector_size,
-            tp_size=tp_size,
-            tp_rank=tp_rank,
-            w3_w1_weight=backend.w3_w1_weight,
-            w2_weight=backend.w2_weight,
-            w3_w1_weight_scale=backend.w3_w1_weight_scale,
-            w2_weight_scale=backend.w2_weight_scale,
-        )
+        loader_module = backend
+        if hasattr(backend, "scaling_vector_size"):
+            loader_module = SimpleNamespace(
+                expert_size_per_partition=backend.expert_size_per_partition,
+                initial_local_expert_ids=backend.initial_local_expert_ids,
+                scaling_vector_size=backend.scaling_vector_size,
+                tp_size=tp_size,
+                tp_rank=tp_rank,
+                w3_w1_weight=backend.w3_w1_weight,
+                w2_weight=backend.w2_weight,
+                w3_w1_weight_scale=backend.w3_w1_weight_scale,
+                w2_weight_scale=backend.w2_weight_scale,
+            )
         for expert_id, tensors in enumerate(bank):
             backend.quant_method.load_packed_mxfp4_expert(
-                proxy,
+                loader_module,
                 global_expert_id=expert_id,
                 local_slot_id=expert_id,
                 w1_weight=tensors["w1"],

From d420dbe350bbfefca462828e5a795b3293f1df5e Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Mon, 3 Aug 2026 19:00:20 -0700
Subject: [PATCH 05/15] [None][test] Move Kimi MoE imports to module scope

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 .../modules/moe/test_kimi_k3_situ_moe.py      | 26 +++++--------------
 1 file changed, 7 insertions(+), 19 deletions(-)

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 27c6044852eb..3703ea4881f4 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
@@ -29,6 +29,12 @@
 import torch.distributed as dist
 from utils.util import check_accuracy
 
+from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_kimi_linear import (
+    _K3_MOE_EP_ENV,
+    _K3_MOE_TP_ENV,
+    KimiK3MoERuntime,
+)
 from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory
 from tensorrt_llm._torch.modules.kimi_k3_moe import KimiK3SparseMoeBlock
 from tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels import (
@@ -38,6 +44,7 @@
 )
 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.mapping import Mapping
 
 situ_supported = pytest.mark.skipif(
     not is_native_situ_supported(),
@@ -420,8 +427,6 @@ def test_fused_forward_without_weights_raises():
 
 
 def test_mapping_records_moe_tp_ep_user_specified():
-    from tensorrt_llm.mapping import Mapping
-
     # Auto default: -1 sentinels resolve to (moe_tp=tp, moe_ep=1) but must
     # NOT be flagged as a user request.
     auto = Mapping(world_size=8, tp_size=8)
@@ -438,13 +443,6 @@ def test_mapping_records_moe_tp_ep_user_specified():
 
 
 def test_kimi_k3_moe_split_selection(monkeypatch):
-    from tensorrt_llm._torch.models.modeling_kimi_linear import (
-        _K3_MOE_EP_ENV,
-        _K3_MOE_TP_ENV,
-        KimiK3MoERuntime,
-    )
-    from tensorrt_llm.mapping import Mapping
-
     monkeypatch.delenv(_K3_MOE_TP_ENV, raising=False)
     monkeypatch.delenv(_K3_MOE_EP_ENV, raising=False)
 
@@ -468,10 +466,6 @@ def test_kimi_k3_moe_split_selection(monkeypatch):
 
 
 def test_kimi_k3_routed_config_preserves_explicit_megamoe_backend():
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoERuntime
-    from tensorrt_llm.mapping import Mapping
-
     model_config = ModelConfig(
         mapping=Mapping(world_size=1, rank=0, tp_size=1),
         moe_backend="MEGAMOE_DEEPGEMM",
@@ -484,10 +478,6 @@ def test_kimi_k3_routed_config_preserves_explicit_megamoe_backend():
 
 
 def test_kimi_k3_routed_config_keeps_trtllm_default():
-    from tensorrt_llm._torch.model_config import ModelConfig
-    from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoERuntime
-    from tensorrt_llm.mapping import Mapping
-
     model_config = ModelConfig(mapping=Mapping(world_size=1, rank=0, tp_size=1))
 
     routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config)
@@ -569,9 +559,7 @@ def _make_routed_moe(
     """Mirror KimiK3MoERuntime's create_moe call on a single-rank mapping."""
     from transformers.configuration_utils import PretrainedConfig
 
-    from tensorrt_llm._torch.model_config import ModelConfig
     from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE, create_moe
-    from tensorrt_llm.mapping import Mapping
     from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig
 
     pretrained_config = PretrainedConfig()

From fb3a9d7de42b3fa0462ec752d85942c257404ee6 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Mon, 3 Aug 2026 19:40:37 -0700
Subject: [PATCH 06/15] [None][fix] honor explicit Kimi MoE backend selection

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 tensorrt_llm/_torch/model_config.py           |  8 ++++++++
 .../_torch/models/modeling_kimi_linear.py     |  8 +-------
 .../modules/moe/test_kimi_k3_situ_moe.py      | 20 +++++++++----------
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index 2d2bc9c314f2..9e0fa9a506ff 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/tensorrt_llm/_torch/model_config.py
@@ -61,6 +61,11 @@
 _DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"}
 _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT = "layers.0.ffn.experts.0.w1.weight"
 
+_KIMI_K3_ARCHITECTURES = {
+    "KimiK3ForConditionalGeneration",
+    "KimiLinearForCausalLM",
+}
+
 _MINIMAX_M3_ARCHITECTURES = {
     "MiniMaxM3SparseForCausalLM",
     "MiniMaxM3SparseForConditionalGeneration",
@@ -367,6 +372,9 @@ def resolve_moe_backend(moe_backend: str,
         if moe_backend.upper() != "AUTO":
             return moe_backend
 
+        if architecture in _KIMI_K3_ARCHITECTURES:
+            return "TRTLLM"
+
         if architecture in _DEEPSEEK_V4_ARCHITECTURES:
             sm_version = get_sm_version()
             if 100 <= sm_version < 120:
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index fe504e30d23f..a20c23e4f984 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -928,13 +928,7 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig:
         routed_model_config._frozen = False
         routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs)
         routed_model_config.mapping = routed_mapping
-        # Preserve K3's TRTLLM-Gen default while allowing the explicit
-        # DeepGEMM MegaMoE opt-in used by the SM100 FP8xFP4 SiTU path.
-        routed_model_config.moe_backend = (
-            "MEGAMOE_DEEPGEMM"
-            if model_config.moe_backend.upper() == "MEGAMOE_DEEPGEMM"
-            else "TRTLLM"
-        )
+        routed_model_config.moe_backend = model_config.moe_backend
         routed_model_config._frozen = True
         return routed_model_config
 
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 3703ea4881f4..037091223ae3 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
@@ -465,24 +465,24 @@ def test_kimi_k3_moe_split_selection(monkeypatch):
     assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2)
 
 
-def test_kimi_k3_routed_config_preserves_explicit_megamoe_backend():
+@pytest.mark.parametrize("backend", ["CUTLASS", "TRTLLM", "MEGAMOE_DEEPGEMM"])
+def test_kimi_k3_routed_config_preserves_explicit_backend(backend):
     model_config = ModelConfig(
         mapping=Mapping(world_size=1, rank=0, tp_size=1),
-        moe_backend="MEGAMOE_DEEPGEMM",
+        moe_backend=backend,
     )
 
     routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config)
 
-    assert routed_model_config.moe_backend == "MEGAMOE_DEEPGEMM"
-    assert model_config.moe_backend == "MEGAMOE_DEEPGEMM"
+    assert routed_model_config.moe_backend == backend
+    assert model_config.moe_backend == backend
 
 
-def test_kimi_k3_routed_config_keeps_trtllm_default():
-    model_config = ModelConfig(mapping=Mapping(world_size=1, rank=0, tp_size=1))
-
-    routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config)
-
-    assert routed_model_config.moe_backend == "TRTLLM"
+@pytest.mark.parametrize(
+    "architecture", ["KimiK3ForConditionalGeneration", "KimiLinearForCausalLM"]
+)
+def test_kimi_k3_moe_auto_backend_defaults_to_trtllm(architecture):
+    assert ModelConfig.resolve_moe_backend("AUTO", architecture) == "TRTLLM"
 
 
 # ---------------------------------------------------------------------------

From 3ed8bd9fbcf95ec0b43b2ceaa8387084ce4698d4 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 02:00:46 -0700
Subject: [PATCH 07/15] [None][chore] Pin DeepGEMM SiTU support upstream

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 3rdparty/fetch_content.json                      | 4 ++--
 scripts/attribution/data/dependency_metadata.yml | 4 ++--
 scripts/attribution/data/files_to_dependency.yml | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json
index 420f56269d4f..85a153113a25 100644
--- a/3rdparty/fetch_content.json
+++ b/3rdparty/fetch_content.json
@@ -31,8 +31,8 @@
     },
     {
       "name": "deepgemm",
-      "git_repository": "https://github.com/longlee0622/DeepGEMM",
-      "git_tag": "c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6",
+      "git_repository": "https://github.com/deepseek-ai/DeepGEMM",
+      "git_tag": "8b1392b978f5a03c828dd1711090d7fb50958b8a",
       "git_submodules_recurse": true,
       "source_subdir": "dont-add-this-project-with-add-subdirectory"
     },
diff --git a/scripts/attribution/data/dependency_metadata.yml b/scripts/attribution/data/dependency_metadata.yml
index 46e990cbff78..ef044adffc11 100644
--- a/scripts/attribution/data/dependency_metadata.yml
+++ b/scripts/attribution/data/dependency_metadata.yml
@@ -23,9 +23,9 @@ cutlass/v4.3.0:
 deep_ep/5be51b228a7c82dbdb213ea58e77bffd12b38af8:
   license: 452b3ef002dc6ec283fb723f0dd84997
   source: https://github.com/deepseek-ai/DeepEP/tree/5be51b228a7c82dbdb213ea58e77bffd12b38af8
-deepgemm/c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6:
+deepgemm/8b1392b978f5a03c828dd1711090d7fb50958b8a:
   license: 452b3ef002dc6ec283fb723f0dd84997
-  source: https://github.com/longlee0622/DeepGEMM/tree/c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6
+  source: https://github.com/deepseek-ai/DeepGEMM/tree/8b1392b978f5a03c828dd1711090d7fb50958b8a
 dlpack/v1.0:
   license: cd9881918c97ec7b4962691660bb733e
 flashmla/1408756a88e52a25196b759eaf8db89d2b51b5a1:
diff --git a/scripts/attribution/data/files_to_dependency.yml b/scripts/attribution/data/files_to_dependency.yml
index 69505bab1ca9..b81c4fbd7df2 100644
--- a/scripts/attribution/data/files_to_dependency.yml
+++ b/scripts/attribution/data/files_to_dependency.yml
@@ -2093,7 +2093,7 @@ deep_ep/5be51b228a7c82dbdb213ea58e77bffd12b38af8:
 - dbc9ea8cf83b20e6ce4c6f383b700f29
 - e4024308dd534f83de92752e1d7cd9a8
 - f41ae95dbbafe6107dd98bf66af018ea
-deepgemm/c35d0ffe906ffcb7382d10a4c497a499ef7ed2d6:
+deepgemm/8b1392b978f5a03c828dd1711090d7fb50958b8a:
 - 07488ed395a262f652b63d0b1c1bb3a8
 - 1101099cb0a9c0489f9e4e49a719941a
 - 111cfba37978c3bad26617cb42407970

From 189c7158e32d87c6e87a9e24c406a2a12004764d Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 02:18:32 -0700
Subject: [PATCH 08/15] [None][chore] Refresh DeepGEMM attribution checksum

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 scripts/attribution/data/files_to_dependency.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/attribution/data/files_to_dependency.yml b/scripts/attribution/data/files_to_dependency.yml
index b81c4fbd7df2..59682bd47f17 100644
--- a/scripts/attribution/data/files_to_dependency.yml
+++ b/scripts/attribution/data/files_to_dependency.yml
@@ -2133,7 +2133,7 @@ deepgemm/8b1392b978f5a03c828dd1711090d7fb50958b8a:
 - 95b8e1175d420f5a05bc9fb0d0c55214
 - 9bf6bbf8da71d31836279a87ffafe108
 - 9e16e23b6894db65b377f93673dc733f
-- 62fe0935dc226557c42dbacf44b2dec0
+- 13ab873c8d6a0b7b2857bbfee317f0a4
 - a3645795bcf4bad7c333094975c558b6
 - a775f6a60d47cd428f4c04289d4e8cb3
 - a9f87d66fb89c05e1ed20a9459c54f68

From 6e454115e3d21472cc8bf5b55384843c38c08798 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 02:23:37 -0700
Subject: [PATCH 09/15] [TRTLLM-15284][fix] Restrict Kimi fused latent path
 backend

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_kimi_linear.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index a20c23e4f984..d56126760ae1 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -87,7 +87,7 @@
 from ..attention_backend import AttentionMetadata
 from ..distributed import AllReduce, AllReduceStrategy
 from ..model_config import ModelConfig
-from ..modules.fused_moe import ConfigurableMoE, create_moe
+from ..modules.fused_moe import ConfigurableMoE, TRTLLMGenFusedMoE, create_moe
 from ..modules.kimi_k3_moe._mlp import KimiK3MLP, KimiK3RMSNorm
 from ..modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate
 from ..modules.linear import Linear as TrtllmLinear
@@ -813,7 +813,7 @@ def __init__(
             not _K3_DISABLE_MIN_LATENCY_LATENT_PROJ
             and not _K3_DISABLE_FUSED_LATENT_DOWN_MXFP8
             and routed_comm is None
-            and hasattr(self.routed_experts.backend, "op_backend")
+            and isinstance(self.routed_experts.backend, TRTLLMGenFusedMoE)
         )
 
         # Shared experts stay replicated (DeepSeek's attention-DP

From a26f590613424b7ee44edf463c8b7e89707b93b5 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 02:30:48 -0700
Subject: [PATCH 10/15] [TRTLLM-15284][fix] Reject unsupported Kimi MoE
 backends

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_kimi_linear.py |  7 +++++++
 .../_torch/modules/moe/test_kimi_k3_situ_moe.py    | 14 +++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index d56126760ae1..4180fc9d3db8 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -891,6 +891,13 @@ def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]:
     def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig:
         """Build a private routed-expert mapping without mutating the shared
         config. Default split is EP-only; see ``_select_moe_tp_ep``."""
+        supported_backends = {"TRTLLM", "MEGAMOE_DEEPGEMM"}
+        if model_config.moe_backend not in supported_backends:
+            raise ValueError(
+                "Kimi K3 SiTU routed experts only support the TRTLLM and "
+                "MEGAMOE_DEEPGEMM backends; "
+                f"got {model_config.moe_backend!r}."
+            )
         if model_config.moe_load_balancer is not None:
             raise NotImplementedError(
                 "Kimi K3 packed-checkpoint streaming does not yet support "
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 037091223ae3..e9ad5a01bb87 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
@@ -465,7 +465,7 @@ def test_kimi_k3_moe_split_selection(monkeypatch):
     assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2)
 
 
-@pytest.mark.parametrize("backend", ["CUTLASS", "TRTLLM", "MEGAMOE_DEEPGEMM"])
+@pytest.mark.parametrize("backend", ["TRTLLM", "MEGAMOE_DEEPGEMM"])
 def test_kimi_k3_routed_config_preserves_explicit_backend(backend):
     model_config = ModelConfig(
         mapping=Mapping(world_size=1, rank=0, tp_size=1),
@@ -478,6 +478,18 @@ def test_kimi_k3_routed_config_preserves_explicit_backend(backend):
     assert model_config.moe_backend == backend
 
 
+def test_kimi_k3_routed_config_rejects_backend_without_situ_support():
+    model_config = ModelConfig(
+        mapping=Mapping(world_size=1, rank=0, tp_size=1),
+        moe_backend="CUTLASS",
+    )
+
+    with pytest.raises(ValueError, match="SiTU routed experts only support"):
+        KimiK3MoERuntime._routed_moe_model_config(model_config)
+
+    assert model_config.moe_backend == "CUTLASS"
+
+
 @pytest.mark.parametrize(
     "architecture", ["KimiK3ForConditionalGeneration", "KimiLinearForCausalLM"]
 )

From 5784902c370f264de2356ec6448027493b87bde7 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 02:36:45 -0700
Subject: [PATCH 11/15] [TRTLLM-15284][fix] Clean up Kimi MegaMoE test group

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 .../modules/moe/test_kimi_k3_situ_moe.py      | 36 ++++++++++++-------
 1 file changed, 23 insertions(+), 13 deletions(-)

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 e9ad5a01bb87..b4eea74ef174 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
@@ -44,6 +44,7 @@
 )
 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
 
 situ_supported = pytest.mark.skipif(
@@ -52,6 +53,25 @@
 )
 
 
+@pytest.fixture
+def _single_rank_nccl_process_group(monkeypatch):
+    if dist.is_initialized():
+        yield
+        return
+    monkeypatch.setenv("MASTER_ADDR", "127.0.0.1")
+    monkeypatch.setenv("MASTER_PORT", str(get_free_port()))
+    monkeypatch.setenv("RANK", "0")
+    monkeypatch.setenv("WORLD_SIZE", "1")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    torch.cuda.set_device(0)
+    dist.init_process_group(backend="nccl", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
 @dataclasses.dataclass
 class _K3Config:
     """Minimal config carrying the fields KimiK3SparseMoeBlock reads."""
@@ -750,7 +770,9 @@ def test_tp8_sharded_forward_matches_whole_expert(num_tokens):
         pytest.param(32, 16, id="experts32-top16"),
     ],
 )
-def test_megamoe_deepgemm_situ_matches_trtllm_gen(num_tokens, num_experts, top_k):
+def test_megamoe_deepgemm_situ_matches_trtllm_gen(
+    num_tokens, num_experts, top_k, _single_rank_nccl_process_group
+):
     """Compare SiTU kernels with identical packed MXFP4 weights and routing.
 
     MegaMoE folds routing weights into the FC1 activation before its MXFP8
@@ -758,18 +780,6 @@ def test_megamoe_deepgemm_situ_matches_trtllm_gen(num_tokens, num_experts, top_k
     quantized graphs therefore need semantic, rather than elementwise,
     parity: high cosine similarity and bounded relative L2 error.
     """
-    if not dist.is_initialized():
-        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
-        os.environ.setdefault("MASTER_PORT", "29561")
-        os.environ.setdefault("RANK", "0")
-        os.environ.setdefault("WORLD_SIZE", "1")
-        os.environ.setdefault("LOCAL_RANK", "0")
-        torch.cuda.set_device(0)
-        dist.init_process_group(
-            backend="nccl",
-            rank=0,
-            world_size=1,
-        )
     bank = _make_packed_expert_bank(num_experts, _TP_INTERMEDIATE, _TP_HIDDEN)
     gate = _make_test_gate(num_experts=num_experts, top_k=top_k)
 

From 60226828c38579a9ae78055b282514d5f21df5fa Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 02:49:34 -0700
Subject: [PATCH 12/15] [TRTLLM-15284][fix] Propagate Kimi SiTU config
 explicitly

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_kimi_linear.py          | 6 ++++++
 .../_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py  | 5 +++++
 tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py  | 6 ++++++
 tests/unittest/_torch/modules/moe/test_moe_backend.py       | 6 ++++--
 4 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 4180fc9d3db8..27fa8530a464 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -778,6 +778,12 @@ def __init__(
                     situ_linear_beta if situ_linear_beta is not None else 1.0
                 ),
             )
+        elif routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM":
+            routed_moe_kwargs.update(
+                activation="situ",
+                situ_beta=float(situ_beta),
+                situ_linear_beta=float(situ_linear_beta if situ_linear_beta is not None else 1.0),
+            )
         self.routed_experts = create_moe(**routed_moe_kwargs)
         if not isinstance(self.routed_experts, ConfigurableMoE):
             raise RuntimeError(
diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
index 132d4ccd3f26..3c6a190d158d 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
@@ -367,8 +367,13 @@ def _resolve_activation_config(
         situ_linear_beta: Optional[float],
     ) -> Tuple[str, Optional[float], Optional[float]]:
         pretrained_config = model_config.pretrained_config
+        text_config = getattr(pretrained_config, "text_config", None)
         config_situ_beta = getattr(pretrained_config, "activation_situ_beta", None)
         config_situ_linear_beta = getattr(pretrained_config, "activation_situ_linear_beta", None)
+        if config_situ_beta is None:
+            config_situ_beta = getattr(text_config, "activation_situ_beta", None)
+        if config_situ_linear_beta is None:
+            config_situ_linear_beta = getattr(text_config, "activation_situ_linear_beta", None)
         if activation is None:
             activation = "situ" if config_situ_beta is not None else "swiglu"
         activation = activation.lower()
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 b4eea74ef174..fc9e8c114943 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
@@ -624,6 +624,12 @@ def _make_routed_moe(
             trtllm_gen_activation_alpha=4.0,
             trtllm_gen_activation_beta=25.0,
         )
+    elif moe_backend == "MEGAMOE_DEEPGEMM":
+        moe_kwargs.update(
+            activation="situ",
+            situ_beta=4.0,
+            situ_linear_beta=25.0,
+        )
     moe = create_moe(**moe_kwargs).cuda()
     assert isinstance(moe, ConfigurableMoE)
     return moe
diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py
index 361f1cb7eed7..2d9840e993f6 100644
--- a/tests/unittest/_torch/modules/moe/test_moe_backend.py
+++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py
@@ -470,8 +470,10 @@ def test_megamoe_deepgemm_cache_derived_state_allocates_symm_buffer():
 def test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config():
     model_config = ModelConfig(
         pretrained_config=SimpleNamespace(
-            activation_situ_beta=4.0,
-            activation_situ_linear_beta=25.0,
+            text_config=SimpleNamespace(
+                activation_situ_beta=4.0,
+                activation_situ_linear_beta=25.0,
+            )
         )
     )
 

From c976d606cd78e1dd741bb808e14730cab434e1b4 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 03:01:49 -0700
Subject: [PATCH 13/15] [TRTLLM-15284][fix] Scope Kimi MegaMoE token capacity

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 examples/kimi_k3/eval_extra_llm_options.yaml  |  4 ++-
 .../_torch/models/modeling_kimi_linear.py     | 10 +++++++
 .../modules/moe/test_kimi_k3_situ_moe.py      | 27 +++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/examples/kimi_k3/eval_extra_llm_options.yaml b/examples/kimi_k3/eval_extra_llm_options.yaml
index c93e2f9607b1..d464e70f4c56 100644
--- a/examples/kimi_k3/eval_extra_llm_options.yaml
+++ b/examples/kimi_k3/eval_extra_llm_options.yaml
@@ -11,7 +11,9 @@ cuda_graph_config:
   enable_padding: true
   max_batch_size: 32
 moe_config:
-  max_num_tokens: 131072
+  # TRTLLM chunking bound. Kimi's MegaMoE path privately raises this to
+  # max_num_tokens * dp_size for per-rank SymmBuffer capacity.
+  max_num_tokens: 33024
   use_low_precision_moe_combine: true
 kv_cache_config:
   enable_block_reuse: false
diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
index 27fa8530a464..9f90bfc07b77 100644
--- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py
+++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py
@@ -942,6 +942,16 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig:
         routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs)
         routed_model_config.mapping = routed_mapping
         routed_model_config.moe_backend = model_config.moe_backend
+        # MegaMoE uses this value as global DP SymmBuffer capacity, then
+        # divides it by EP size for the per-rank allocation. Other backends
+        # keep the user-configured value as their MoE chunking bound.
+        # Preserve an explicitly larger capacity.
+        if routed_model_config.moe_backend == "MEGAMOE_DEEPGEMM":
+            default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size
+            routed_model_config.moe_max_num_tokens = max(
+                int(routed_model_config.moe_max_num_tokens or 0),
+                default_moe_max_num_tokens,
+            )
         routed_model_config._frozen = True
         return routed_model_config
 
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 fc9e8c114943..0bef2f2cdaa2 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
@@ -498,6 +498,33 @@ def test_kimi_k3_routed_config_preserves_explicit_backend(backend):
     assert model_config.moe_backend == backend
 
 
+@pytest.mark.parametrize(
+    "backend,expected_moe_max_num_tokens",
+    [
+        pytest.param("TRTLLM", 33024, id="trtllm"),
+        pytest.param("MEGAMOE_DEEPGEMM", 131072, id="megamoe"),
+    ],
+)
+def test_kimi_k3_routed_config_scopes_megamoe_capacity(backend, expected_moe_max_num_tokens):
+    model_config = ModelConfig(
+        mapping=Mapping(
+            world_size=16,
+            rank=0,
+            tp_size=16,
+            moe_ep_size=16,
+            enable_attention_dp=True,
+        ),
+        max_num_tokens=8192,
+        moe_max_num_tokens=33024,
+        moe_backend=backend,
+    )
+
+    routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config)
+
+    assert routed_model_config.moe_max_num_tokens == expected_moe_max_num_tokens
+    assert model_config.moe_max_num_tokens == 33024
+
+
 def test_kimi_k3_routed_config_rejects_backend_without_situ_support():
     model_config = ModelConfig(
         mapping=Mapping(world_size=1, rank=0, tp_size=1),

From 7ffa0483b60aa1793febf5a93f59668dca009f42 Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 03:50:14 -0700
Subject: [PATCH 14/15] [TRTLLM-15284][test] Fix Kimi MegaMoE test setup

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py | 6 ------
 1 file changed, 6 deletions(-)

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 0bef2f2cdaa2..fa1c554df489 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
@@ -651,12 +651,6 @@ def _make_routed_moe(
             trtllm_gen_activation_alpha=4.0,
             trtllm_gen_activation_beta=25.0,
         )
-    elif moe_backend == "MEGAMOE_DEEPGEMM":
-        moe_kwargs.update(
-            activation="situ",
-            situ_beta=4.0,
-            situ_linear_beta=25.0,
-        )
     moe = create_moe(**moe_kwargs).cuda()
     assert isinstance(moe, ConfigurableMoE)
     return moe

From 936648edf54825777708a0fa91af2a9db0604eac Mon Sep 17 00:00:00 2001
From: Jonas Li <6110159+longlee0622@users.noreply.github.com>
Date: Tue, 11 Aug 2026 17:30:53 -0700
Subject: [PATCH 15/15] [TRTLLM-15284][test] Clear MegaMoE cache in test
 teardown

Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com>
---
 tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py | 4 ++++
 1 file changed, 4 insertions(+)

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 fa1c554df489..e17497be77f4 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
@@ -36,6 +36,9 @@
     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 import KimiK3SparseMoeBlock
 from tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels import (
     is_native_situ_supported,
@@ -68,6 +71,7 @@ def _single_rank_nccl_process_group(monkeypatch):
     try:
         yield
     finally:
+        _MEGA_MOE_SYMM_BUFFER_CACHE.clear()
         if dist.is_initialized():
             dist.destroy_process_group()