[TRTLLM-15284][feat] add Kimi K3 SiTU MegaMoE support

3rdparty/fetch_content.json

@@ -32,7 +32,7 @@
     {
       "name": "deepgemm",
       "git_repository": "https://github.com/deepseek-ai/DeepGEMM",
-      "git_tag": "f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca",
+      "git_tag": "8b1392b978f5a03c828dd1711090d7fb50958b8a",
       "git_submodules_recurse": true,
       "source_subdir": "dont-add-this-project-with-add-subdirectory"
     },

examples/kimi_k3/eval_extra_llm_options.yaml

@@ -11,6 +11,8 @@ cuda_graph_config:
   enable_padding: true
   max_batch_size: 32
 moe_config:
+  # 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:

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/8b1392b978f5a03c828dd1711090d7fb50958b8a:
   license: 452b3ef002dc6ec283fb723f0dd84997
-  source: https://github.com/deepseek-ai/DeepGEMM/tree/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca
+  source: https://github.com/deepseek-ai/DeepGEMM/tree/8b1392b978f5a03c828dd1711090d7fb50958b8a
 dlpack/v1.0:
   license: cd9881918c97ec7b4962691660bb733e
 flashmla/1408756a88e52a25196b759eaf8db89d2b51b5a1:

scripts/attribution/data/files_to_dependency.yml

@@ -2093,7 +2093,7 @@ deep_ep/5be51b228a7c82dbdb213ea58e77bffd12b38af8:
 - dbc9ea8cf83b20e6ce4c6f383b700f29
 - e4024308dd534f83de92752e1d7cd9a8
 - f41ae95dbbafe6107dd98bf66af018ea
-deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca:
+deepgemm/8b1392b978f5a03c828dd1711090d7fb50958b8a:
 - 07488ed395a262f652b63d0b1c1bb3a8
 - 1101099cb0a9c0489f9e4e49a719941a
 - 111cfba37978c3bad26617cb42407970
@@ -2133,7 +2133,7 @@ deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca:
 - 95b8e1175d420f5a05bc9fb0d0c55214
 - 9bf6bbf8da71d31836279a87ffafe108
 - 9e16e23b6894db65b377f93673dc733f
-- a359f41b0ec5d67c1d9fda93dc65c5b5
+- 13ab873c8d6a0b7b2857bbfee317f0a4
 - a3645795bcf4bad7c333094975c558b6
 - a775f6a60d47cd428f4c04289d4e8cb3
 - a9f87d66fb89c05e1ed20a9459c54f68
@@ -2157,7 +2157,7 @@ deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca:
 - ea07df16c9a083277f55ae219c85f39d
 - ef5544cabdf0490063f2b2959f62a8cf
 - ef712bc72e01afbfdbceef82c6b49174
-- f43084b87a14bcf56920e1e3b1ad2dba
+- cfcffb4170384e7699c65ed604fbd1d8
 - f5a3009221d096c818cf26fbdb4d9693
 - f5d68cc5860baa1de18f4aee4ccb0cc2
 - f7f27b18dae31db0aad429bfaaf4615e

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:

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
@@ -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,38 @@ 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
+                ),
+            )
+        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(
                 "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 +819,7 @@ def __init__(
             not _K3_DISABLE_MIN_LATENCY_LATENT_PROJ
             and not _K3_DISABLE_FUSED_LATENT_DOWN_MXFP8
             and routed_comm is None
+            and isinstance(self.routed_experts.backend, TRTLLMGenFusedMoE)
         )
 
         # Shared experts stay replicated (DeepSeek's attention-DP
@@ -878,6 +897,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 "
@@ -915,7 +941,17 @@ 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"
+        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
 

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,46 @@ 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
+        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()
+        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 +762,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)