[TRTLLM-14959][refactor] declare MoE backend capabilities instead of checking exact classes

tensorrt_llm/_torch/models/modeling_deepseekv4.py

@@ -74,16 +74,13 @@
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
 from ..moe.fused_moe import (
-    CutlassFusedMoE,
+    DEFAULT_MOE_ACTIVATION,
     DeepSeekV4MoeRoutingMethod,
     MoEWeightLoadingMode,
-    TritonFusedMoE,
-    TRTLLMGenFusedMoE,
+    SwigluActivation,
     create_moe,
     is_moe_weight_owner,
-    resolve_moe_cls,
 )
-from ..moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE
 from ..peft.lora.layer import LoraLayer
 from ..speculative import SpecMetadata, get_num_extra_kv_tokens
 from ..utils import (
@@ -1545,61 +1542,15 @@ def __init__(
         if override_quant_config is not None and experts_quant_config is model_config.quant_config:
             experts_quant_config = override_quant_config
 
+        # One uniform clamp for the whole layer; the selected backend's
+        # ``activation_support`` decides whether its kernels take that as a
+        # per-expert ``float32`` buffer or a baked scalar.
         swiglu_limit = getattr(config, "swiglu_limit", None)
-        moe_swiglu_limit = None
-        if swiglu_limit is not None:
-            # `create_moe` only accepts swiglu_limit for these MoE classes;
-            # ask the resolver rather than the backend string so that a
-            # degradation (e.g. TRTLLM/CUTEDSL/DENSEGEMM dropping back to
-            # CutlassFusedMoE on unsupported quant) is accounted for here too.
-            moe_cls = resolve_moe_cls(
-                model_config,
-                override_quant_config=experts_quant_config,
-                dtype=dtype,
-                # Same routing object as create_moe below.
-                routing=self.gate.routing_method,
-                # create_moe below passes no bias and no swiglu alpha/beta, so
-                # it resolves with the plain SwiGLU package. Say so here too:
-                # leaving this unknown lets gates abstain that create_moe
-                # rejects, and the two calls would pick different backends.
-                swiglu_gptoss_style=False,
-                layer_idx=layer_idx,
-            )
-            supports_swiglu_limit = moe_cls in (
-                CutlassFusedMoE,
-                TritonFusedMoE,
-                TRTLLMGenFusedMoE,
-                DeepGemmFusedMoE,
-            )
-            # NVFP4 routed-expert path: the TRTLLM-Gen fp4-block-scale fused-MoE
-            # cubin produces near-zero accuracy without bias even when
-            # swiglu_limit is supplied; drop the limit there until the cubin
-            # gains a no-bias clamp variant. MXFP4 variants are unaffected.
-            kernel_requires_bias_for_swiglu_limit = (
-                moe_cls is TRTLLMGenFusedMoE and experts_quant_config.quant_mode.has_nvfp4()
-            )
-            # DeepSeek-V4 supplies a uniform scalar limit. The TRTLLM-Gen FP8
-            # path consumes it directly and rejects the redundant tensor.
-            requires_scalar_only_swiglu_limit = (
-                moe_cls is TRTLLMGenFusedMoE
-                and experts_quant_config.quant_mode.has_fp8_block_scales()
-            )
-            if (
-                supports_swiglu_limit
-                and not kernel_requires_bias_for_swiglu_limit
-                and not requires_scalar_only_swiglu_limit
-            ):
-                moe_load_balancer_config = getattr(model_config, "moe_load_balancer", None)
-                num_slots = (
-                    moe_load_balancer_config.num_slots
-                    if moe_load_balancer_config and moe_load_balancer_config.num_slots
-                    else num_experts
-                )
-                local_num_slots = num_slots // model_config.mapping.moe_ep_size
-                device = "cuda" if torch.cuda.is_available() else "cpu"
-                moe_swiglu_limit = torch.full(
-                    (local_num_slots,), float(swiglu_limit), dtype=torch.float32, device=device
-                )
+        moe_activation = (
+            SwigluActivation(clamp=float(swiglu_limit))
+            if swiglu_limit is not None
+            else DEFAULT_MOE_ACTIVATION
+        )
 
         self.experts = create_moe(
             num_experts=num_experts,
@@ -1619,8 +1570,7 @@ def __init__(
                 if experts_quant_config.layer_quant_mode.is_int4_weight_only_per_group()
                 else MoEWeightLoadingMode.VANILLA
             ),
-            swiglu_limit=moe_swiglu_limit,
-            swiglu_limit_scalar=(float(swiglu_limit) if swiglu_limit is not None else None),
+            activation=moe_activation,
         )
 
         self.mapping = model_config.mapping

tensorrt_llm/_torch/models/modeling_gemma4.py

@@ -27,6 +27,7 @@
 
 from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper
 from tensorrt_llm._torch.modules.qk_norm_attention import QKNormRoPEAttention
+from tensorrt_llm._torch.moe.fused_moe.activation import SimpleActivation
 from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe
 from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode
 from tensorrt_llm._torch.moe.fused_moe.routing import BaseMoeRoutingMethod
@@ -669,7 +670,7 @@ def __init__(
             reduce_results=True,
             model_config=model_config,
             layer_idx=layer_idx,
-            activation_type=ActivationType.Geglu,
+            activation=SimpleActivation(kind=ActivationType.Geglu),
             # VANILLA mode: preprocess_weights splits 3D gate_up_proj into per-expert w1/w3
             weight_loading_mode=MoEWeightLoadingMode.VANILLA,
         )

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -25,7 +25,8 @@
 # isort and yapf will fight against each other here, so we disable isort
 # isort: off
 from ..moe.fused_moe import (MoEWeightLoadingMode, RenormalizeMoeRoutingMethod,
-                             TritonFusedMoE, create_moe, is_moe_weight_owner)
+                             SwigluBiasActivation, TritonFusedMoE, create_moe,
+                             is_moe_weight_owner)
 # isort: on
 from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
@@ -172,15 +173,12 @@ def __init__(
             output_dtype=torch.bfloat16
             if config.moe_backend.upper() == "TRTLLM" else torch.float32)
 
-        self.swiglu_alpha = torch.tensor(
-            [1.702] * (self.num_slots // config.mapping.moe_ep_size),
-            dtype=torch.float32).cuda()
-        self.swiglu_beta = torch.tensor(
-            [1.0] * (self.num_slots // config.mapping.moe_ep_size),
-            dtype=torch.float32).cuda()
-        self.swiglu_limit = torch.tensor(
-            [7.0] * (self.num_slots // config.mapping.moe_ep_size),
-            dtype=torch.float32).cuda()
+        # gpt-oss constants are uniform across experts; the backend broadcasts
+        # them to whatever per-expert shape its kernels index, sized by the
+        # local slot count it resolved.
+        self.moe_activation = SwigluBiasActivation(gate_sigmoid_scale=1.702,
+                                                   linear_offset=1.0,
+                                                   clamp=7.0)
         # Prepare MoE creation parameters
         moe_params = {
             'routing_method': self.routing_method,
@@ -192,9 +190,7 @@ def __init__(
             'model_config': config,
             'weight_loading_mode': MoEWeightLoadingMode.FUSED_GATE_UP_PROJ,
             'bias': True,
-            'swiglu_alpha': self.swiglu_alpha,
-            'swiglu_beta': self.swiglu_beta,
-            'swiglu_limit': self.swiglu_limit,
+            'activation': self.moe_activation,
             'layer_idx': self.layer_idx,
         }
 

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -117,10 +117,8 @@
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
 from ..modules.situ import SituAndMul
-from ..moe.fused_moe import ConfigurableMoE, create_moe
-from ..moe.fused_moe.interface import _compute_ep_partition
+from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, create_moe
 from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod
-from ..utils import ActivationType, ActType_TrtllmGen
 from .modeling_speculative import SpecDecOneEngineForCausalLM
 from .modeling_utils import DecoderModel, register_auto_model, run_concurrently
 
@@ -1120,6 +1118,16 @@ def __init__(
             layer_idx=layer_idx,
             # Let CommunicationFactory select the best available strategy.
             communication_method=None,
+            activation=SiTuActivation(
+                gate_softcap=situ_beta,
+                linear_softcap=situ_linear_beta,
+            ),
+            # A MegaMoE request that silently degraded to CUTLASS would be
+            # benchmarked as if it were MegaMoE, and the decline is easy to
+            # trigger (EP-only, own token / top-k limits). Fail in the resolver
+            # instead, which reports the rejection trail.
+            allow_backend_degradation=routed_moe_model_config.moe_backend
+            not in ("MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"),
         )
         # trtllm-gen ships SiTu cubins for exactly one dtype combination
         # (``Bmm_MxE4m3_MxE2m1MxE4m3`` = MXFP8 act x MXFP4 weight) and has no
@@ -1141,82 +1149,11 @@ def __init__(
                 "CUTLASS or MEGAMOE_CUTEDSL."
             )
 
-        if routed_moe_model_config.moe_backend == "CUTLASS":
-            # Size the per-expert SiTU constants with the same ceil/floor
-            # partition the MoE backend uses for ``expert_size_per_partition``.
-            # A plain ``num_experts // ep_size`` is one element short on the
-            # first ``num_experts % ep_size`` ranks, which trips the
-            # ``swiglu_alpha must have num_experts_on_rank elements`` check in
-            # moeOp.cpp. K3's 384 experts divide evenly at EP8/EP16, so the
-            # mismatch is latent there but real for any uneven split.
-            local_num_experts, _, _ = _compute_ep_partition(
-                self.num_experts,
-                routed_moe_model_config.mapping.moe_ep_size,
-                routed_moe_model_config.mapping.moe_ep_rank,
-            )
-            device = torch.device("cuda", torch.cuda.current_device())
-            self.routed_situ_alpha = torch.full(
-                (local_num_experts,), float(situ_beta), dtype=torch.float32, device=device
-            )
-            self.routed_situ_beta = torch.full(
-                (local_num_experts,),
-                situ_linear_beta,
-                dtype=torch.float32,
-                device=device,
-            )
-            routed_moe_kwargs.update(
-                activation_type=ActivationType.SiTu,
-                swiglu_alpha=self.routed_situ_alpha,
-                swiglu_beta=self.routed_situ_beta,
-            )
-        elif 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=situ_beta,
-                trtllm_gen_activation_beta=situ_linear_beta,
-            )
-        elif routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM":
-            routed_moe_kwargs.update(
-                activation="situ",
-                situ_beta=situ_beta,
-                situ_linear_beta=situ_linear_beta,
-            )
-        # MEGAMOE_CUTEDSL has no branch on purpose. ``MegaMoECuteDsl`` resolves
-        # SiTU from the pretrained config in ``_resolve_activation_config``
-        # (``activation=None`` -> "situ" when ``activation_situ_beta`` is
-        # present), and ``create_moe`` currently rejects the explicit
-        # ``activation``/``situ_beta``/``situ_linear_beta`` trio for anything
-        # other than ``MegaMoEDeepGemm``, so passing them here would raise.
-        # Unifying that plumbing is tracked in TRTLLM-15649.
         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 ..moe.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 routed_moe_model_config.moe_backend == "MEGAMOE_CUTEDSL":
-            from ..moe.fused_moe.mega_moe import MegaMoECuteDsl
-
-            # Same guard as MEGAMOE_DEEPGEMM above, and for the same reason:
-            # create_moe silently falls back when a backend declines the
-            # config, and for MegaMoE the decline is easy to trigger (it is
-            # EP-only and has its own token/top-k limits), so an explicit
-            # request that quietly became CUTLASS would be measured as if it
-            # were MegaMoE.
-            if not isinstance(self.routed_experts.backend, MegaMoECuteDsl):
-                raise RuntimeError(
-                    "Kimi K3 explicitly requested MEGAMOE_CUTEDSL, 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 "

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -62,15 +62,9 @@
 )
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
-from ..moe.fused_moe import MiniMaxM3MoeRoutingMethod, create_moe
+from ..moe.fused_moe import MiniMaxM3MoeRoutingMethod, SwigluBiasActivation, create_moe
 from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph
-from ..utils import (
-    ActivationType,
-    AuxStreamType,
-    EventType,
-    get_model_extra_attrs,
-    is_torch_compiling,
-)
+from ..utils import AuxStreamType, EventType, get_model_extra_attrs, is_torch_compiling
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper
 from .modeling_utils import (
@@ -300,33 +294,6 @@ def _build_swiglu_oai_dense_mlp(
     )
 
 
-def _resolve_minimax_m3_expert_size_per_partition(
-    num_experts: int,
-    mapping: Mapping,
-    moe_load_balancer_config,
-) -> int:
-    """Compute the local expert/slot count the MoE module will resolve.
-
-    Sizes the per-expert SwiGLU parameter tensors we hand to
-    ``create_moe`` to match what the backend sees. Some backends assert
-    ``swiglu_alpha.shape == (expert_size_per_partition,)`` at construct
-    time; a mismatch surfaces as an opaque CUDA-side shape failure.
-
-    Priority:
-      1. EPLB config → ``num_slots // moe_ep_size``.
-      2. Plain EP    → ``num_experts // moe_ep_size`` (with ``max(1, ...)``
-                        for tiny test configs).
-    """
-    ep_size = mapping.moe_ep_size
-    if moe_load_balancer_config is not None and moe_load_balancer_config.num_slots:
-        # Mirror ``MoeLoadBalancerConfig.num_local_slots`` without
-        # requiring ``.setup(ep_rank, ep_size)`` to have been called
-        # (the ``num_local_slots`` property raises otherwise).
-        return moe_load_balancer_config.num_slots // ep_size
-
-    return max(1, num_experts // ep_size)
-
-
 class MiniMaxM3Gate(nn.Module):
     """MiniMax-M3 router gate: float32 sigmoid scoring + per-expert bias correction.
 
@@ -457,31 +424,14 @@ def __init__(
         self.swiglu_beta_value = 1.0  # SGLang's ``(up + 1)`` offset in swiglu_no_interleaved.
         self.swiglu_limit_value = float(getattr(config, "swiglu_limit", 7.0))
 
-        # Size the per-expert SwiGLU parameter tensors from the same
-        # ``expert_size_per_partition`` the MoE module will resolve, not
-        # from a hand-rolled ``num_slots // ep_size`` guess. EPLB and
-        # DWDP can shift the local slot count (see
-        # :func:`_resolve_minimax_m3_expert_size_per_partition` for the
-        # priority order), and some backends assert
-        # ``swiglu_alpha.shape == (expert_size_per_partition,)``.
-        moe_load_balancer_config = model_config.moe_load_balancer
-        self.expert_size_per_partition = _resolve_minimax_m3_expert_size_per_partition(
-            num_experts=self.num_experts,
-            mapping=model_config.mapping,
-            moe_load_balancer_config=moe_load_balancer_config,
+        # One value per layer, not per expert: the MoE backend broadcasts these
+        # to whatever per-expert shape its kernels index (or bakes them as
+        # scalars), sized by the slot count it actually resolved.
+        self.moe_activation = SwigluBiasActivation(
+            gate_sigmoid_scale=self.swiglu_alpha_value,
+            linear_offset=self.swiglu_beta_value,
+            clamp=self.swiglu_limit_value,
         )
-        self.swiglu_alpha = torch.tensor(
-            [self.swiglu_alpha_value] * self.expert_size_per_partition,
-            dtype=torch.float32,
-        ).cuda()
-        self.swiglu_beta = torch.tensor(
-            [self.swiglu_beta_value] * self.expert_size_per_partition,
-            dtype=torch.float32,
-        ).cuda()
-        self.swiglu_limit = torch.tensor(
-            [self.swiglu_limit_value] * self.expert_size_per_partition,
-            dtype=torch.float32,
-        ).cuda()
 
         # Router gate owns the float32 projection weight, the per-expert
         # ``e_score_correction_bias``, and the ``routing_method`` it
@@ -508,24 +458,8 @@ def __init__(
             model_config=model_config,
             layer_idx=layer_idx,
             override_quant_config=experts_quant_config,
-            swiglu_alpha=self.swiglu_alpha,
-            swiglu_beta=self.swiglu_beta,
-            swiglu_limit=self.swiglu_limit,
-            activation_type=ActivationType.SwigluBias,
+            activation=self.moe_activation,
         )
-        # Defensive: if a future MoE-resolution path (new load-balancer
-        # mode, new DWDP variant) shifts the local expert count in a
-        # way our resolver doesn't yet model, fail here with a
-        # diagnostic message instead of inside a CUDA-side shape
-        # assertion deep in a backend kernel dispatch.
-        resolved = getattr(self.experts, "expert_size_per_partition", None)
-        assert resolved is None or resolved == self.expert_size_per_partition, (
-            f"MiniMax-M3 SwiGLU sizing mismatch: pre-create_moe estimate "
-            f"{self.expert_size_per_partition} != MoE-resolved {resolved}. "
-            f"Update _resolve_minimax_m3_expert_size_per_partition to "
-            f"match the MoE module's resolved layout."
-        )
-
         # Shared expert: dense MLP fused into MoE output. Constructed
         # with ``is_shared_expert=True`` so ``reduce_output=False``
         # (the external AllReduce below performs the combined