[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

@@ -662,7 +662,7 @@ def _convert_mla_projections_to_fp8_weight_read(model: nn.Module) -> int:
 
 
 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,
@@ -706,7 +706,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,
@@ -716,20 +716,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 "
@@ -820,6 +838,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 "
@@ -857,7 +882,29 @@ 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
+            configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0)
+            if configured_moe_max_num_tokens < default_moe_max_num_tokens:
+                logger.info_once(
+                    "Kimi K3 MegaMoE raises moe_max_num_tokens from "
+                    f"{configured_moe_max_num_tokens} to {default_moe_max_num_tokens} "
+                    "because the global DP SymmBuffer requires capacity for "
+                    "max_num_tokens * dp_size.",
+                    key=(
+                        "kimi_k3_megamoe_capacity_override_"
+                        f"{configured_moe_max_num_tokens}_{default_moe_max_num_tokens}"
+                    ),
+                )
+            routed_model_config.moe_max_num_tokens = max(
+                configured_moe_max_num_tokens,
+                default_moe_max_num_tokens,
+            )
         routed_model_config._frozen = True
         return routed_model_config
 

tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py

@@ -162,6 +162,9 @@ def __init__(
         apply_router_weight_on_input: bool = False,
         layer_idx: Optional[int] = None,
         override_quant_config: Optional["QuantConfig"] = None,
+        activation: Optional[str] = None,
+        situ_beta: Optional[float] = None,
+        situ_linear_beta: Optional[float] = None,
         trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None,
         trtllm_gen_activation_alpha: Optional[float] = None,
         trtllm_gen_activation_beta: Optional[float] = None,
@@ -197,6 +200,9 @@ def __init__(
             model_config=model_config,
             routing_method=routing_method,
             override_quant_config=override_quant_config,
+            activation=activation,
+            situ_beta=situ_beta,
+            situ_linear_beta=situ_linear_beta,
             trtllm_gen_activation_type=trtllm_gen_activation_type,
             trtllm_gen_activation_alpha=trtllm_gen_activation_alpha,
             trtllm_gen_activation_beta=trtllm_gen_activation_beta,
@@ -283,6 +289,9 @@ def _create_and_sync_backend(
         model_config: ModelConfig,
         routing_method: BaseMoeRoutingMethod,
         override_quant_config: Optional["QuantConfig"],
+        activation: Optional[str],
+        situ_beta: Optional[float],
+        situ_linear_beta: Optional[float],
         trtllm_gen_activation_type: Optional[ActType_TrtllmGen],
         trtllm_gen_activation_alpha: Optional[float],
         trtllm_gen_activation_beta: Optional[float],
@@ -340,6 +349,9 @@ def _create_and_sync_backend(
                 swiglu_limit_scalar=kwargs.get("swiglu_limit_scalar"),
                 init_load_balancer=False,
                 activation_type=self.activation_type,
+                activation=activation,
+                situ_beta=situ_beta,
+                situ_linear_beta=situ_linear_beta,
                 trtllm_gen_activation_type=trtllm_gen_activation_type,
                 trtllm_gen_activation_alpha=trtllm_gen_activation_alpha,
                 trtllm_gen_activation_beta=trtllm_gen_activation_beta,

tensorrt_llm/_torch/modules/fused_moe/create_moe.py

@@ -289,6 +289,9 @@ def create_moe_backend(
     swiglu_limit_scalar: Optional[float] = None,
     init_load_balancer: bool = True,
     activation_type: ActivationType = ActivationType.Swiglu,
+    activation: Optional[str] = None,
+    situ_beta: Optional[float] = None,
+    situ_linear_beta: Optional[float] = None,
     trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None,
     trtllm_gen_activation_alpha: Optional[float] = None,
     trtllm_gen_activation_beta: Optional[float] = None,
@@ -315,6 +318,9 @@ def create_moe_backend(
         swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4)
         swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8)
         activation_type: Activation type
+        activation: Optional MegaMoE DeepGEMM activation name
+        situ_beta: Optional MegaMoE DeepGEMM SiTU beta
+        situ_linear_beta: Optional MegaMoE DeepGEMM SiTU linear beta
         trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type
         trtllm_gen_activation_alpha: Optional backend-local activation alpha
         trtllm_gen_activation_beta: Optional backend-local activation beta
@@ -410,6 +416,12 @@ def create_moe_backend(
             trtllm_gen_activation_beta=trtllm_gen_activation_beta,
         )
 
+    if any(value is not None
+           for value in (activation, situ_beta,
+                         situ_linear_beta)) and moe_cls is not MegaMoEDeepGemm:
+        raise ValueError("MegaMoE DeepGEMM activation options require "
+                         f"MegaMoEDeepGemm, got {moe_cls.__name__}")
+
     if any(value is not None for value in (trtllm_gen_activation_type,
                                            trtllm_gen_activation_alpha,
                                            trtllm_gen_activation_beta)):
@@ -554,7 +566,12 @@ def create_moe_backend(
                                               if swiglu_limit is not None else
                                               swiglu_limit_scalar)
         else:
-            megamoe_kwargs["swiglu_limit_scalar"] = swiglu_limit_scalar
+            megamoe_kwargs.update(
+                swiglu_limit_scalar=swiglu_limit_scalar,
+                activation=activation,
+                situ_beta=situ_beta,
+                situ_linear_beta=situ_linear_beta,
+            )
         return moe_cls(**megamoe_kwargs)
     else:
         raise ValueError(f"Unsupported moe backend: {moe_cls}")
@@ -579,6 +596,9 @@ def create_moe(
     swiglu_limit: Optional[torch.Tensor] = None,
     swiglu_limit_scalar: Optional[float] = None,
     activation_type: ActivationType = ActivationType.Swiglu,
+    activation: Optional[str] = None,
+    situ_beta: Optional[float] = None,
+    situ_linear_beta: Optional[float] = None,
     trtllm_gen_activation_type: Optional[ActType_TrtllmGen] = None,
     trtllm_gen_activation_alpha: Optional[float] = None,
     trtllm_gen_activation_beta: Optional[float] = None,
@@ -606,6 +626,9 @@ def create_moe(
         swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4)
         swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8)
         activation_type: Activation type
+        activation: Optional MegaMoE DeepGEMM activation name
+        situ_beta: Optional MegaMoE DeepGEMM SiTU beta
+        situ_linear_beta: Optional MegaMoE DeepGEMM SiTU linear beta
         trtllm_gen_activation_type: Optional TRTLLM-Gen backend-local activation type
         trtllm_gen_activation_alpha: Optional backend-local activation alpha
         trtllm_gen_activation_beta: Optional backend-local activation beta
@@ -635,6 +658,13 @@ def create_moe(
 
     moe_cls = resolve_moe_cls(model_config, routing_method, dtype,
                               override_quant_config, layer_idx)
+    if (any(value is not None
+            for value in (activation, situ_beta, situ_linear_beta))
+            and moe_cls is not MegaMoEDeepGemm):
+        raise ValueError(
+            "MegaMoE DeepGEMM activation options require "
+            "MegaMoEDeepGemm without backend fallback, but resolved "
+            f"{moe_cls.__name__}.")
     if (any(value is not None for value in (trtllm_gen_activation_type,
                                             trtllm_gen_activation_alpha,
                                             trtllm_gen_activation_beta))
@@ -666,6 +696,9 @@ def create_moe(
             swiglu_limit=swiglu_limit,
             swiglu_limit_scalar=swiglu_limit_scalar,
             activation_type=activation_type,
+            activation=activation,
+            situ_beta=situ_beta,
+            situ_linear_beta=situ_linear_beta,
             trtllm_gen_activation_type=trtllm_gen_activation_type,
             trtllm_gen_activation_alpha=trtllm_gen_activation_alpha,
             trtllm_gen_activation_beta=trtllm_gen_activation_beta,
@@ -695,6 +728,9 @@ def create_moe(
         swiglu_limit=swiglu_limit,
         swiglu_limit_scalar=swiglu_limit_scalar,
         activation_type=activation_type,
+        activation=activation,
+        situ_beta=situ_beta,
+        situ_linear_beta=situ_linear_beta,
         trtllm_gen_activation_type=trtllm_gen_activation_type,
         trtllm_gen_activation_alpha=trtllm_gen_activation_alpha,
         trtllm_gen_activation_beta=trtllm_gen_activation_beta,