From 622e70b1a3d6d8f7ab6f7e04fe8b75c2ff39bfe0 Mon Sep 17 00:00:00 2001 From: guqiqi <29116997+guqiqi@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:52:08 +0000 Subject: [PATCH 1/4] [None][feat] Add qwen3_8 / kimi_k3 bench_moe presets and activation plumbing Signed-off-by: guqiqi <29116997+guqiqi@users.noreply.github.com> --- .../modules/fused_moe/fused_moe_deepgemm.py | 8 +++- .../bench_moe/BENCH_MOE_USER_GUIDE.md | 8 ++++ tests/microbenchmarks/bench_moe/build.py | 3 ++ tests/microbenchmarks/bench_moe/cli.py | 10 +++++ tests/microbenchmarks/bench_moe/search.py | 12 +++++- tests/microbenchmarks/bench_moe/specs.py | 37 +++++++++++++++++++ 6 files changed, 76 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py index 39b713613fce..76685dde6de2 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/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() diff --git a/tests/microbenchmarks/bench_moe/BENCH_MOE_USER_GUIDE.md b/tests/microbenchmarks/bench_moe/BENCH_MOE_USER_GUIDE.md index 5748a9dd8aa5..661d84a3086b 100644 --- a/tests/microbenchmarks/bench_moe/BENCH_MOE_USER_GUIDE.md +++ b/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 diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index 63949d74d701..e734dd370f32 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -220,6 +220,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 +265,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( @@ -285,6 +287,7 @@ 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, ) if quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8: diff --git a/tests/microbenchmarks/bench_moe/cli.py b/tests/microbenchmarks/bench_moe/cli.py index 0b516983ada1..0c9815cf0907 100644 --- a/tests/microbenchmarks/bench_moe/cli.py +++ b/tests/microbenchmarks/bench_moe/cli.py @@ -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,6 +554,7 @@ 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. @@ -576,6 +585,7 @@ def _resolve_model_from_args(args: argparse.Namespace) -> ModelSpec: swiglu_alpha=base.swiglu_alpha, swiglu_beta=base.swiglu_beta, swiglu_limit=base.swiglu_limit, + activation_type=args.activation or base.activation_type, ) diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py index 6aec8ec4fed8..e6fa7e74a994 100644 --- a/tests/microbenchmarks/bench_moe/search.py +++ b/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 diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py index be8f24bc649f..616db594f7a3 100644 --- a/tests/microbenchmarks/bench_moe/specs.py +++ b/tests/microbenchmarks/bench_moe/specs.py @@ -35,6 +35,7 @@ RenormalizeNaiveMoeRoutingMethod, SigmoidRenormMoeRoutingMethod, ) +from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm.models.modeling_utils import QuantAlgo from .backend import MoeBackendType, MoeModelConfig @@ -49,6 +50,11 @@ "SIGMOID_RENORM": SigmoidRenormMoeRoutingMethod, } +_ACTIVATIONS: Dict[str, ActivationType] = { + "SWIGLU": ActivationType.Swiglu, + "RELU2": ActivationType.Relu2, +} + _ROUTING_NAME_BY_CLS: Dict[type, str] = {cls: name for name, cls in _ROUTING_METHODS.items()} _ALL_BACKENDS: List[str] = [b.value for b in MoeBackendType] _COMM_METHODS: Tuple[str, ...] = ( @@ -117,6 +123,7 @@ class ModelSpec: swiglu_alpha: float = 1.0 swiglu_beta: float = 0.0 swiglu_limit: float = float("inf") + activation_type: str = "SWIGLU" def __post_init__(self) -> None: if self.n_shared_experts < 0: @@ -130,6 +137,10 @@ def __post_init__(self) -> None: def routing_method_cls(self) -> type: return _ROUTING_METHODS[self.routing_method] + @property + def activation_type_enum(self) -> ActivationType: + return _ACTIVATIONS[self.activation_type] + @property def quant_algo_enum(self) -> Optional[QuantAlgo]: return QuantAlgo[self.quant_algo] if self.quant_algo is not None else None @@ -386,4 +397,30 @@ class RunResult: swiglu_beta=1.0, swiglu_limit=7.0, ), + # Qwen3.8, 2.4T total / A95B activated + "qwen3_8": ModelSpec( + name="qwen3_8", + num_experts=512, + top_k=10, + hidden_size=8192, + intermediate_size=2048, + quant_algo=None, + routing_method="RENORMALIZE", + ), + # Kimi-K3 + # Activation is "situ" (situ(g) * up, situ(g) = b*tanh(g/b)*sigmoid(g)), which + # TensorRT-LLM does not implement here. situ is *gated*, so SWIGLU has the + # same fc1 fan-out, GEMM shapes and comm volume -- a faithful perf proxy, + # differing only in the epilogue elementwise math. + "kimi_k3": ModelSpec( + name="kimi_k3", + num_experts=896, + top_k=16, + hidden_size=3584, + intermediate_size=3072, + quant_algo=None, + routing_method="DEEPSEEK_V3", + n_group=1, + topk_group=1, + ), } From 49b140697984c9c50af482fbc153f190fd0cff37 Mon Sep 17 00:00:00 2001 From: guqiqi <29116997+guqiqi@users.noreply.github.com> Date: Thu, 20 Aug 2026 02:52:09 +0000 Subject: [PATCH 2/4] [None][feat] Run kimi_k3 with its SiTU epilogue on TRTLLM and MegaMoE DeepGEMM Signed-off-by: guqiqi <29116997+guqiqi@users.noreply.github.com> --- tests/microbenchmarks/bench_moe/build.py | 47 +++++++++++++++++++ .../microbenchmarks/bench_moe/case_runner.py | 2 + tests/microbenchmarks/bench_moe/cli.py | 46 +++++++----------- tests/microbenchmarks/bench_moe/results.py | 1 + tests/microbenchmarks/bench_moe/specs.py | 19 ++++++++ 5 files changed, 85 insertions(+), 30 deletions(-) diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index e734dd370f32..9d1437f8ce99 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -30,6 +30,7 @@ import torch from tensorrt_llm._torch.modules.fused_moe.interface import MoESchedulerKind, MoEWeightLoadingMode +from tensorrt_llm._torch.utils import ActType_TrtllmGen from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo @@ -83,6 +84,22 @@ 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. The two backends expose it differently: + MegaMoE stores the resolved name, TRTLLM-Gen a predicate. + """ + backend = getattr(moe, "backend", None) or moe + if getattr(backend, "activation", None) == "situ" or getattr( + backend, "is_situ_activation", False + ): + 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 +114,35 @@ 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]) -> 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. Both paths also hard-require W4A8_MXFP4_MXFP8, so 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 or quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8: + return {} + backend = moe_backend.upper() + if backend == "MEGAMOE_DEEPGEMM": + return { + "activation": "situ", + "situ_beta": model.situ_beta, + "situ_linear_beta": model.situ_linear_beta, + } + if backend == "TRTLLM": + # 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, + } + 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 @@ -288,6 +334,7 @@ def _build_moe_module( 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, + **_situ_kwargs(model, moe_backend, quant_algo), ) if quant_algo == QuantAlgo.W4A8_MXFP4_MXFP8: diff --git a/tests/microbenchmarks/bench_moe/case_runner.py b/tests/microbenchmarks/bench_moe/case_runner.py index cea0008ab67e..0cfbcf411133 100644 --- a/tests/microbenchmarks/bench_moe/case_runner.py +++ b/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(): diff --git a/tests/microbenchmarks/bench_moe/cli.py b/tests/microbenchmarks/bench_moe/cli.py index 0c9815cf0907..ccaa00e203bf 100644 --- a/tests/microbenchmarks/bench_moe/cli.py +++ b/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 @@ -557,36 +557,22 @@ def _resolve_model_from_args(args: argparse.Namespace) -> ModelSpec: 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, - activation_type=args.activation or base.activation_type, - ) + 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]: diff --git a/tests/microbenchmarks/bench_moe/results.py b/tests/microbenchmarks/bench_moe/results.py index c24f7bbd39c1..cf4f38a2ee88 100644 --- a/tests/microbenchmarks/bench_moe/results.py +++ b/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, diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py index 616db594f7a3..0abb5dfd2eff 100644 --- a/tests/microbenchmarks/bench_moe/specs.py +++ b/tests/microbenchmarks/bench_moe/specs.py @@ -124,6 +124,11 @@ class ModelSpec: swiglu_beta: float = 0.0 swiglu_limit: float = float("inf") activation_type: str = "SWIGLU" + # SiTU constants (Kimi-K3: activation_situ_beta / activation_situ_linear_beta). + # SiTU is gated, so ``activation_type`` stays SWIGLU and only the epilogue + # changes; backends take it through their own parameters (see build.py). + situ_beta: Optional[float] = None + situ_linear_beta: Optional[float] = None def __post_init__(self) -> None: if self.n_shared_experts < 0: @@ -132,6 +137,15 @@ def __post_init__(self) -> None: raise ValueError( f"shared_expert_mode must be 'fused' or 'unfused', got {self.shared_expert_mode!r}." ) + if (self.situ_beta is None) != (self.situ_linear_beta is None): + raise ValueError("situ_beta and situ_linear_beta must be set together.") + if self.situ_beta is not None: + if self.situ_beta <= 0 or self.situ_linear_beta <= 0: + raise ValueError("SiTU betas must be positive.") + if self.activation_type != "SWIGLU": + raise ValueError("SiTU requires the SwiGLU fc1 geometry (activation_type=SWIGLU).") + if self.swiglu_limit != float("inf"): + raise ValueError("SiTU is mutually exclusive with a gate/up clamp.") @property def routing_method_cls(self) -> type: @@ -275,6 +289,9 @@ class RunResult: actual_backend: Optional[str] = None actual_comm_method: Optional[str] = None actual_comm_fallback_reason: Optional[str] = None + # Epilogue the built module ran ("swiglu" / "situ"), read back rather + # than echoed from the spec. + actual_epilogue_activation: Optional[str] = None scheduler_kind: Optional[str] = None moe_ep_size: Optional[int] = None moe_tp_size: Optional[int] = None @@ -422,5 +439,7 @@ class RunResult: routing_method="DEEPSEEK_V3", n_group=1, topk_group=1, + situ_beta=4.0, # text_config.activation_situ_beta + situ_linear_beta=25.0, # text_config.activation_situ_linear_beta ), } From 07f9fcc6642c81871f6f48e5425dc8192eab9ec2 Mon Sep 17 00:00:00 2001 From: guqiqi <29116997+guqiqi@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:01:32 +0000 Subject: [PATCH 3/4] [None][feat] Pass SiTU kwargs for MegaMoE CuteDSL NVFP4 in bench_moe Wire kimi_k3 SiTU through MEGAMOE_CUTEDSL+NVFP4 the same way as DeepGEMM, instead of dropping it behind the W4A8-only gate. Signed-off-by: guqiqi <29116997+guqiqi@users.noreply.github.com> --- tests/microbenchmarks/bench_moe/build.py | 31 +++++++++++++----------- tests/microbenchmarks/bench_moe/specs.py | 9 +++---- 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index 9d1437f8ce99..cb61deb45333 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -89,8 +89,8 @@ def _epilogue_activation_name(moe) -> str: 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. The two backends expose it differently: - MegaMoE stores the resolved name, TRTLLM-Gen a predicate. + reporting what was asked for. MegaMoE DeepGEMM / CuteDSL store the + resolved name; TRTLLM-Gen exposes a predicate. """ backend = getattr(moe, "backend", None) or moe if getattr(backend, "activation", None) == "situ" or getattr( @@ -119,21 +119,24 @@ def _situ_kwargs(model: ModelSpec, moe_backend: str, quant_algo: Optional[QuantA 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. Both paths also hard-require W4A8_MXFP4_MXFP8, so 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. + 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 or quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8: + if model.situ_beta is None: return {} backend = moe_backend.upper() - if backend == "MEGAMOE_DEEPGEMM": - return { - "activation": "situ", - "situ_beta": model.situ_beta, - "situ_linear_beta": model.situ_linear_beta, - } - if backend == "TRTLLM": + 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, diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py index 0abb5dfd2eff..641ba6debdc8 100644 --- a/tests/microbenchmarks/bench_moe/specs.py +++ b/tests/microbenchmarks/bench_moe/specs.py @@ -424,11 +424,10 @@ class RunResult: quant_algo=None, routing_method="RENORMALIZE", ), - # Kimi-K3 - # Activation is "situ" (situ(g) * up, situ(g) = b*tanh(g/b)*sigmoid(g)), which - # TensorRT-LLM does not implement here. situ is *gated*, so SWIGLU has the - # same fc1 fan-out, GEMM shapes and comm volume -- a faithful perf proxy, - # differing only in the epilogue elementwise math. + # Kimi-K3. SiTU is gated, so ``activation_type`` stays SWIGLU (same fc1 + # geometry). ``build._situ_kwargs`` switches the epilogue for the backends + # that implement it (TRTLLM / MEGAMOE_DEEPGEMM on W4A8_MXFP4_MXFP8, + # MEGAMOE_CUTEDSL on NVFP4); other combinations keep the SwiGLU proxy. "kimi_k3": ModelSpec( name="kimi_k3", num_experts=896, From 6ce07ad4211905fe7ba3b321ad6853c98d2fc9c8 Mon Sep 17 00:00:00 2001 From: kikig <29116997+guqiqi@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:27:06 -0700 Subject: [PATCH 4/4] [None][feat] Wire CUTLASS NVFP4 SiTU into bench_moe for kimi_k3. bench_moe already switched MegaMoE CuteDSL NVFP4 to SiTU; CUTLASS uses ActivationType.SiTu plus per-rank alpha/beta, so pass those through for the same quant. Signed-off-by: kikig <29116997+guqiqi@users.noreply.github.com> --- tests/microbenchmarks/bench_moe/build.py | 50 +++++++++++++++++++++--- tests/microbenchmarks/bench_moe/specs.py | 3 +- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index cb61deb45333..dd1c7e3057be 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -29,8 +29,12 @@ import torch -from tensorrt_llm._torch.modules.fused_moe.interface import MoESchedulerKind, MoEWeightLoadingMode -from tensorrt_llm._torch.utils import ActType_TrtllmGen +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 @@ -90,13 +94,17 @@ def _epilogue_activation_name(moe) -> str: 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. + 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" @@ -114,7 +122,12 @@ 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]) -> Dict: +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`` @@ -143,6 +156,29 @@ def _situ_kwargs(model: ModelSpec, moe_backend: str, quant_algo: Optional[QuantA "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 {} @@ -323,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, @@ -337,8 +374,9 @@ def _build_moe_module( 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, - **_situ_kwargs(model, moe_backend, quant_algo), ) + 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( diff --git a/tests/microbenchmarks/bench_moe/specs.py b/tests/microbenchmarks/bench_moe/specs.py index 641ba6debdc8..b62ac964905f 100644 --- a/tests/microbenchmarks/bench_moe/specs.py +++ b/tests/microbenchmarks/bench_moe/specs.py @@ -427,7 +427,8 @@ class RunResult: # Kimi-K3. SiTU is gated, so ``activation_type`` stays SWIGLU (same fc1 # geometry). ``build._situ_kwargs`` switches the epilogue for the backends # that implement it (TRTLLM / MEGAMOE_DEEPGEMM on W4A8_MXFP4_MXFP8, - # MEGAMOE_CUTEDSL on NVFP4); other combinations keep the SwiGLU proxy. + # MEGAMOE_CUTEDSL / CUTLASS on NVFP4); other combinations keep the SwiGLU + # proxy. "kimi_k3": ModelSpec( name="kimi_k3", num_experts=896,