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..dd1c7e3057be 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/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( 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 0b516983ada1..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 @@ -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]: 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/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..b62ac964905f 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,12 @@ class ModelSpec: swiglu_alpha: float = 1.0 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: @@ -125,11 +137,24 @@ 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: 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 @@ -264,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 @@ -386,4 +414,32 @@ 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. 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 / CUTLASS on NVFP4); other combinations keep the SwiGLU + # proxy. + "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, + situ_beta=4.0, # text_config.activation_situ_beta + situ_linear_beta=25.0, # text_config.activation_situ_linear_beta + ), }