[None][feat] Add qwen3_8 / kimi_k3 bench_moe presets and activation plumbing

tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py

@@ -27,7 +27,7 @@
 
 from ...memory_buffer_utils import get_memory_buffers
 from ...model_config import ModelConfig
-from ...utils import AuxStreamType, Fp4QuantizedTensor
+from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor
 from .fused_moe_cutlass import CutlassFusedMoE
 from .impl_contract import (MoEDeployment, MoEEligibility, MoEInputRequirement,
                             MoEProblem, MoERejectReason, MoERunContext,
@@ -812,6 +812,12 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility:
                 "DeepGemmFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)"
             )
 
+        # DeepGemmFusedMoE The inter-GEMM activation is a hardcoded silu_and_mul
+        if p.activation_type != ActivationType.Swiglu:
+            return _reject(
+                MoERejectReason.ACTIVATION_UNSUPPORTED,
+                f"DeepGemmFusedMoE only supports SwiGLU (got {p.activation})")
+
         # Only FP8_BLOCK_SCALES is supported
         if quant_algo == QuantAlgo.FP8_BLOCK_SCALES:
             return MoEEligibility.ok()

tests/microbenchmarks/bench_moe/BENCH_MOE_USER_GUIDE.md

@@ -663,11 +663,19 @@ includes full requested and observed dispatch/expert matrices.
 | `deepseek_v3` | 256 | 8 | 7168 | 2048 | `FP8_BLOCK_SCALES` | `DEEPSEEK_V3` |
 | `deepseek_r1` | 256 | 8 | 7168 | 2048 | `FP8_BLOCK_SCALES` | `DEEPSEEK_V3` |
 | `kimi_k2` | 384 | 8 | 7168 | 2048 | `FP8_BLOCK_SCALES` | `DEEPSEEK_V3` |
+| `kimi_k3` | 896 | 16 | 3584 † | 3072 | pass `--quant` | `DEEPSEEK_V3` |
 | `glm_5` | 256 | 8 | 6144 | 2048 | pass `--quant` | `DEEPSEEK_V3` |
 | `deepseek_v4_pro` | 384 | 6 | 7168 | 3072 | pass `--quant` | `RENORMALIZE` |
 | `deepseek_v4_flash` | 256 | 6 | 4096 | 2048 | pass `--quant` | `RENORMALIZE` |
+| `qwen3_8` | 512 | 10 | 8192 | 2048 | pass `--quant` | `RENORMALIZE` |
 | `gpt_oss_120b` | 128 | 4 | 2880 | 2880 | `W4A8_MXFP4_MXFP8` | `RENORMALIZE` |
 
+† **Latent MoE.** Some models project tokens down to a smaller latent dimension
+before dispatch and back up after combine, so the routed experts never see the
+model's `hidden_size`. For those, `Hidden` is the latent dimension — that is what
+both the expert GEMMs and the all-to-all payload actually run at. The outer
+down/up projections sit outside the MoE layer and are not benchmarked.
+
 Custom shapes can be used instead of `--model`:
 
 ```bash

tests/microbenchmarks/bench_moe/build.py

@@ -29,7 +29,12 @@
 
 import torch
 
-from tensorrt_llm._torch.modules.fused_moe.interface import MoESchedulerKind, MoEWeightLoadingMode
+from tensorrt_llm._torch.modules.fused_moe.interface import (
+    MoESchedulerKind,
+    MoEWeightLoadingMode,
+    _compute_ep_partition,
+)
+from tensorrt_llm._torch.utils import ActivationType, ActType_TrtllmGen
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantAlgo
 
@@ -83,6 +88,26 @@ def _comm_method_name(moe) -> str:
     return type(comm).__name__
 
 
+def _epilogue_activation_name(moe) -> str:
+    """Return the epilogue the built module actually runs: ``"situ"`` or ``"swiglu"``.
+
+    The SiTU request can be dropped for reasons the spec cannot see (wrong
+    backend, wrong quant, upstream fallback), so read it back rather than
+    reporting what was asked for. MegaMoE DeepGEMM / CuteDSL store the
+    resolved name; TRTLLM-Gen exposes a predicate; CUTLASS uses
+    ``ActivationType.SiTu``.
+    """
+    backend = getattr(moe, "backend", None) or moe
+    if getattr(backend, "activation", None) == "situ" or getattr(
+        backend, "is_situ_activation", False
+    ):
+        return "situ"
+    activation_type = getattr(backend, "activation_type", None)
+    if activation_type is not None and ActivationType(activation_type) == ActivationType.SiTu:
+        return "situ"
+    return "swiglu"
+
+
 def _calculate_num_chunks_safe(moe, all_rank_num_tokens: List[int]) -> Optional[int]:
     """Best-effort lookup of ``num_chunks`` for the case we are about to time."""
     scheduler = getattr(moe, "scheduler", None)
@@ -97,6 +122,66 @@ def _calculate_num_chunks_safe(moe, all_rank_num_tokens: List[int]) -> Optional[
         return None
 
 
+def _situ_kwargs(
+    model: ModelSpec,
+    moe_backend: str,
+    quant_algo: Optional[QuantAlgo],
+    mapping: Optional[Mapping] = None,
+) -> Dict:
+    """``create_moe`` kwargs that switch the epilogue to SiTU, or ``{}``.
+
+    Each backend takes SiTU through its own parameters and ``create_moe``
+    REJECTS them on any other backend, so this dispatches instead of passing
+    one set everywhere. Quant is paired with the backend that actually
+    implements the epilogue; a spec carrying SiTU constants falls back to
+    the SwiGLU proxy elsewhere rather than failing the case -- SiTU is
+    gated, so the GEMM shapes and comm volume are the same either way.
+    """
+    if model.situ_beta is None:
+        return {}
+    backend = moe_backend.upper()
+    mega_situ = {
+        "activation": "situ",
+        "situ_beta": model.situ_beta,
+        "situ_linear_beta": model.situ_linear_beta,
+    }
+    if backend == "MEGAMOE_DEEPGEMM" and quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8:
+        return mega_situ
+    if backend == "MEGAMOE_CUTEDSL" and quant_algo == QuantAlgo.NVFP4:
+        return mega_situ
+    if backend == "TRTLLM" and quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8:
+        # Cubin alpha is the gate-side beta, cubin beta the linear-side one.
+        return {
+            "trtllm_gen_activation_type": ActType_TrtllmGen.SiTu,
+            "trtllm_gen_activation_alpha": model.situ_beta,
+            "trtllm_gen_activation_beta": model.situ_linear_beta,
+        }
+    if backend == "CUTLASS" and quant_algo == QuantAlgo.NVFP4:
+        # CUTLASS takes SiTU as ActivationType plus per-rank alpha/beta
+        # (same packing as modeling_kimi_linear). Size the tensors with
+        # the ceil/floor EP partition so uneven splits stay valid.
+        ep_size = 1 if mapping is None else max(mapping.moe_ep_size, 1)
+        ep_rank = 0 if mapping is None else mapping.moe_ep_rank
+        local_num_experts, _, _ = _compute_ep_partition(model.num_experts, ep_size, ep_rank)
+        device = torch.device("cuda", torch.cuda.current_device())
+        return {
+            "activation_type": ActivationType.SiTu,
+            "swiglu_alpha": torch.full(
+                (local_num_experts,),
+                float(model.situ_beta),
+                dtype=torch.float32,
+                device=device,
+            ),
+            "swiglu_beta": torch.full(
+                (local_num_experts,),
+                float(model.situ_linear_beta),
+                dtype=torch.float32,
+                device=device,
+            ),
+        }
+    return {}
+
+
 def _create_moe_for_benchmark(**kwargs):
     ensure_cute_dsl_importable_for_benchmark()
     from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe
@@ -220,6 +305,7 @@ def _build_moe_module(
 
     mc = model.to_moe_model_config()
     swiglu_gptoss_style = model.swiglu_gptoss_style
+    activation_type = model.activation_type_enum
 
     routing_method = _create_routing_method(
         model.routing_method_cls,
@@ -264,6 +350,7 @@ def _build_moe_module(
         swiglu_beta=model.swiglu_beta if swiglu_gptoss_style else None,
         swiglu_limit=model.swiglu_limit if swiglu_gptoss_style else None,
         num_local_experts=num_local_experts,
+        activation_type=activation_type,
     )
 
     weight_loading_mode = getattr(
@@ -272,7 +359,8 @@ def _build_moe_module(
 
     swiglu_tensors = quantize_util.get_swiglu_tensors()
 
-    moe = _create_moe_for_benchmark(
+    # Merge then unpack so SiTU can override activation_type / swiglu_alpha-beta.
+    moe_kwargs = dict(
         routing_method=routing_method,
         num_experts=mc.num_experts,
         hidden_size=mc.hidden_size,
@@ -285,7 +373,10 @@ def _build_moe_module(
         swiglu_alpha=swiglu_tensors["swiglu_alpha"] if swiglu_tensors else None,
         swiglu_beta=swiglu_tensors["swiglu_beta"] if swiglu_tensors else None,
         swiglu_limit=swiglu_tensors["swiglu_limit"] if swiglu_tensors else None,
+        activation_type=activation_type,
     )
+    moe_kwargs.update(_situ_kwargs(model, moe_backend, quant_algo, mapping))
+    moe = _create_moe_for_benchmark(**moe_kwargs)
 
     if quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8:
         weights, _ref_weights, _ref_kwargs = quantize_util.prepare_weights_from_backend(

tests/microbenchmarks/bench_moe/case_runner.py

@@ -35,6 +35,7 @@
     _build_moe_module,
     _calculate_num_chunks_safe,
     _comm_method_name,
+    _epilogue_activation_name,
     _scheduler_kind_name,
 )
 from .mapping import _build_mapping_from_config, _resolve_mapping_layout
@@ -726,6 +727,7 @@ def _run_one_candidate(
         result.actual_backend = _backend_name_from_module(moe)
         result.scheduler_kind = _scheduler_kind_name(moe)
         result.actual_comm_method = _comm_method_name(moe)
+        result.actual_epilogue_activation = _epilogue_activation_name(moe)
         result.num_chunks = _calculate_num_chunks_safe(moe, all_rank_num_tokens)
 
         if result.actual_backend != config.backend.upper():

tests/microbenchmarks/bench_moe/cli.py

@@ -20,7 +20,7 @@
 import argparse
 import json
 import sys
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch
@@ -39,6 +39,7 @@
     _resolve_search_from_args,
 )
 from .specs import (
+    _ACTIVATIONS,
     _ALL_BACKENDS,
     _COMM_METHODS,
     _ROUTING_METHODS,
@@ -230,6 +231,13 @@ def parse_args() -> argparse.Namespace:
             "default; custom shapes must specify an explicit method."
         ),
     )
+    model_group.add_argument(
+        "--activation",
+        type=lambda s: str(s).upper(),
+        default=None,
+        choices=sorted(_ACTIVATIONS),
+        help="Expert activation. Defaults to the model's value, else SWIGLU.",
+    )
 
     workload_group = parser.add_argument_group("Workload shape")
     workload_group.add_argument(
@@ -546,37 +554,25 @@ def _resolve_model_from_args(args: argparse.Namespace) -> ModelSpec:
             topk_group=args.topk_group,
             n_shared_experts=int(args.n_shared_experts) if args.n_shared_experts is not None else 0,
             shared_expert_mode=args.shared_expert_mode,
+            activation_type=args.activation or "SWIGLU",
         )
 
-    # Built-in model with optional per-field overrides.
-    if routing == "AUTO":
-        routing = base.routing_method
-
-    quant_name: Optional[str]
+    # Built-in model: override only what the CLI actually provided, so any
+    # preset field without a flag (swiglu_*, situ_*, ...) is inherited.
+    overrides: Dict[str, Any] = {"shared_expert_mode": args.shared_expert_mode}
+    if routing != "AUTO":
+        overrides["routing_method"] = routing
     if args.quant is not None:
-        quant_name = args.quant.name
-    else:
-        quant_name = base.quant_algo
-    return ModelSpec(
-        name=base.name,
-        num_experts=int(args.num_experts) if args.num_experts is not None else base.num_experts,
-        top_k=int(args.top_k) if args.top_k is not None else base.top_k,
-        hidden_size=int(args.hidden_size) if args.hidden_size is not None else base.hidden_size,
-        intermediate_size=int(args.intermediate_size)
-        if args.intermediate_size is not None
-        else base.intermediate_size,
-        quant_algo=quant_name,
-        routing_method=routing,
-        n_group=args.n_group if args.n_group is not None else base.n_group,
-        topk_group=args.topk_group if args.topk_group is not None else base.topk_group,
-        n_shared_experts=int(args.n_shared_experts)
-        if args.n_shared_experts is not None
-        else base.n_shared_experts,
-        shared_expert_mode=args.shared_expert_mode,
-        swiglu_alpha=base.swiglu_alpha,
-        swiglu_beta=base.swiglu_beta,
-        swiglu_limit=base.swiglu_limit,
-    )
+        overrides["quant_algo"] = args.quant.name
+    if args.activation is not None:
+        overrides["activation_type"] = args.activation
+    for field in ("num_experts", "top_k", "hidden_size", "intermediate_size", "n_shared_experts"):
+        if getattr(args, field) is not None:
+            overrides[field] = int(getattr(args, field))
+    for field in ("n_group", "topk_group"):
+        if getattr(args, field) is not None:
+            overrides[field] = getattr(args, field)
+    return replace(base, **overrides)
 
 
 def _resolve_workloads_from_args(args: argparse.Namespace) -> List[WorkloadSpec]:

tests/microbenchmarks/bench_moe/results.py

@@ -355,6 +355,7 @@ def _runresult_to_row(result: RunResult) -> Dict[str, Any]:
             "backend": result.actual_backend,
             "comm_method": result.actual_comm_method,
             "comm_fallback_reason": result.actual_comm_fallback_reason,
+            "epilogue_activation": result.actual_epilogue_activation,
             "scheduler_kind": result.scheduler_kind,
             "moe_ep_size": result.moe_ep_size,
             "moe_tp_size": result.moe_tp_size,

tests/microbenchmarks/bench_moe/search.py

@@ -27,9 +27,11 @@
 from tensorrt_llm._torch.modules.fused_moe.impl_contract import (
     MoEDeployment,
     MoEProblem,
+    canonical_activation,
     canonical_quant,
 )
 from tensorrt_llm._torch.modules.fused_moe.impl_environment import collect_moe_environment
+from tensorrt_llm._torch.utils import ActivationType
 from tensorrt_llm._utils import local_mpi_size
 from tensorrt_llm.models.modeling_utils import QuantAlgo
 
@@ -63,6 +65,7 @@ def _check_backend_can_implement(
     quant_algo: Optional[QuantAlgo],
     dtype_activation: torch.dtype,
     swiglu_gptoss_style: bool,
+    activation_type: ActivationType,
 ) -> Tuple[bool, Optional[str]]:
     """Resolve backend_str to its MoE class and ask whether it can serve this.
 
@@ -78,6 +81,9 @@ def _check_backend_can_implement(
         quant=canonical_quant(quant_algo),
         dtype_act=dtype_activation,
         swiglu_gptoss_style=swiglu_gptoss_style,
+        # Defaults to SwiGLU when unset, which would make every upstream
+        # activation gate evaluate the wrong activation.
+        activation=canonical_activation(activation_type),
     )
     deployment = MoEDeployment(
         ep_size=1,
@@ -175,7 +181,11 @@ def is_candidate_valid(
     """Return ``(ok, reason)`` based on backend / mapping / comm gates."""
     # Backend can_implement gate.
     ok, reason = _check_backend_can_implement(
-        config.backend, model.quant_algo_enum, act_dtype, model.swiglu_gptoss_style
+        config.backend,
+        model.quant_algo_enum,
+        act_dtype,
+        model.swiglu_gptoss_style,
+        model.activation_type_enum,
     )
     if not ok:
         return False, reason