[None][perf] Triton FP8-activation MoE for decode-sized batches on Hopper

tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md

@@ -332,13 +332,14 @@ parent of every leaf) and re-exports `raise_moe_lora_multichunk_unsupported`.
 | `base.py` | `CutlassFusedMoEBase` — construction, the layer-format check in `create_weights`, and the defaults: a pass-through `quantize_input` and a `run_moe` that calls the op with the leaf's `FUSED_MOE_OP_FLAGS` (the weight lifecycle itself is `MoEWeightOwnerMixin`'s). Also the one `torch.ops.trtllm.fused_moe` call (`_run_fused_moe_op`, which a leaf's `run_moe` calls with the operands the `run_moe` contract does not carry), `FusedMoeOpFlags`, and the `quantize_input` implementations several formats bind (`quantize_input_static_e4m3`, `quantize_input_mxfp8`) |
 | `eligibility.py` | The checks each leaf's `can_implement` composes |
 | `lora.py` | `CutlassMoELoraMixin` — routed-expert LoRA, mixed in only by the two leaves the op can fuse it for |
-| `<format>.py` | One registered leaf each, named after its format: `unquantized.py` for `none`, `fp8_block_scales_hopper.py` / `fp8_block_scales_triton.py` for the two `fp8_block_scales` leaves, the `quant` segment otherwise |
+| `<format>.py` | One registered leaf each, named after its format: `unquantized.py` for `none`, `fp8_block_scales_hopper.py` / `fp8_block_scales_hopper_triton.py` / `fp8_block_scales_triton.py` for the three `fp8_block_scales` leaves, the `quant` segment otherwise |
 
 | `quant` segment | Leaf | Id | What differs from the base |
 |---|---|---|---|
 | `none` | `TrtllmCutlassUnquantizedImpl` | `trtllm.cutlass.grouped_gemm.none` | Fuses routed-expert LoRA |
 | `fp8` | `TrtllmCutlassFp8Impl` | `trtllm.cutlass.grouped_gemm.fp8` | Static per-tensor e4m3 input; fuses routed-expert LoRA |
 | `fp8_block_scales` | `DeepgemmCudaHopperFp8BlockScalesImpl` | `deepgemm.cuda.hopper_grouped_gemm.fp8_block_scales` | SM90; `CutlassMoeFCRunner` hands FC1/FC2 to the vendored DeepGEMM |
+| `fp8_block_scales` | `TrtllmTritonFp8ActGroupedGemmFp8BlockScalesImpl` | `trtllm.triton.fp8_act_grouped_gemm.fp8_block_scales` | SM90, opt-in (`TLLM_MOE_ENABLE_SM90_TRITON_FP8_DECODE=1`, else it declines with `PATH_NOT_ENABLED`); subclasses the DeepGEMM leaf. SwiGLU only, no expert bias. `run_moe` runs Triton FP8-activation GEMMs (`run_triton_fp8_block_scale_moe_fp8_act`) when it receives at most `TRITON_FP8_ACT_MAX_TOKENS` rows after dispatch, no `moe_output` buffer (so never under NVLinkOneSided), BF16 input and a clamp uniform across experts (`_takes_triton_path`); otherwise the parent's fused op |
 | `fp8_block_scales` | `TrtllmTritonFp8BlockScalesImpl` | `trtllm.triton.grouped_gemm.fp8_block_scales` | SM120/121, BF16 activations; a Triton grouped GEMM that never enters C++ (`issues_fused_moe_op = False`, so no build gate); runs SwiGLU/GeGLU only, takes no expert bias and no activation constant but a uniform pre-SiLU SwiGLU clamp |
 | `nvfp4` | `TrtllmCutlassNvfp4Impl` | `trtllm.cutlass.grouped_gemm.nvfp4` | Quantizes activations to NVFP4 itself |
 | `w4a16_nvfp4` | `TrtllmCutlassW4a16Nvfp4Impl` | `trtllm.cutlass.grouped_gemm.w4a16_nvfp4` | Dequantizes the weights every forward, then runs unquantized |
@@ -355,12 +356,14 @@ parent of every leaf) and re-exports `raise_moe_lora_multichunk_unsupported`.
   copies `scheduler_kind`, `capabilities` and `input_requirement` off the
   descriptor, so the contract has one source, and nothing in the base branches
   on `quant_config`.
-- **FP8 block scales splits by SM.** Its two leaves carry the providers and
+- **FP8 block scales splits by SM.** Its three leaves carry the providers and
   techniques of the kernels they run (DeepGEMM, Triton), yet stay in the family
   because the `CUTLASS` literal serves FP8 block scales on SM90 and SM120/121.
+  The two SM90 leaves overlap: the opt-in Triton FP8-activation leaf ranks
+  first and declines unless its `opt_in_flag` is set.
   `find_cutlass_leaf(quant_algo, *, sm=None)` (and `cutlass_leaf`, which
   raises instead of returning `None`) is a lookup, not an eligibility check, and
-  reads the SM only to choose between those two.
+  reads the SM, and a leaf's `opt_in_flag`, only to choose among those leaves.
 
 ### TRTLLM-Gen leaves (`fused_moe/trtllm_gen/`)
 
@@ -1104,7 +1107,8 @@ as `<Provider><Technique><Quant>Impl`. The kernel segment enters when two ids
 differ in that segment alone: the class that would otherwise collide spells it
 out between technique and quant, in id order
 (`DeepgemmCudaHopperFp8BlockScalesImpl` next to
-`DeepgemmCudaFp8BlockScalesImpl`). The `trtllm.cutedsl.*.nvfp4` ids
+`DeepgemmCudaFp8BlockScalesImpl`, `TrtllmTritonFp8ActGroupedGemmFp8BlockScalesImpl`
+next to `TrtllmTritonFp8BlockScalesImpl`). The `trtllm.cutedsl.*.nvfp4` ids
 (`mega_moe`, `fused_fc12`, `dense_gemm`, `grouped_gemm`) differ in
 `kernel_name` alone, so their classes spell the kernel out;
 `TrtllmCutedslGroupedGemmBf16Impl` keeps it to match its NVFP4 sibling. Quant

tensorrt_llm/_torch/moe/fused_moe/cutlass/__init__.py

@@ -25,6 +25,7 @@
 from .base import CutlassFusedMoEBase
 from .fp8 import TrtllmCutlassFp8Impl
 from .fp8_block_scales_hopper import DeepgemmCudaHopperFp8BlockScalesImpl
+from .fp8_block_scales_hopper_triton import TrtllmTritonFp8ActGroupedGemmFp8BlockScalesImpl
 from .fp8_block_scales_triton import TrtllmTritonFp8BlockScalesImpl
 from .lora import CutlassMoELoraMixin, raise_moe_lora_multichunk_unsupported
 from .mxfp8 import TrtllmCutlassMxfp8Impl
@@ -38,10 +39,13 @@
 from .w8a16 import TrtllmCutlassW8a16Impl
 
 #: Every leaf in this family, in the order ``IMPL_PRIORITY`` ranks them and
-#: ``moe_resolution`` falls back to them. The order cannot change which leaf
-#: admits a problem: the ``quant`` segments are disjoint, except for the two
-#: ``fp8_block_scales`` leaves, whose SM sets are.
+#: ``moe_resolution`` falls back to them. The ``quant`` segments are disjoint
+#: except for the three ``fp8_block_scales`` leaves. The SM120/121 one has an SM
+#: set of its own. The two SM90 ones overlap, so the opt-in Triton
+#: FP8-activation leaf comes first: when opted in it wins, and otherwise it
+#: declines and the DeepGEMM leaf serves.
 CUTLASS_LEAVES: tuple[type[CutlassFusedMoEBase], ...] = (
+    TrtllmTritonFp8ActGroupedGemmFp8BlockScalesImpl,
     DeepgemmCudaHopperFp8BlockScalesImpl,
     TrtllmTritonFp8BlockScalesImpl,
     TrtllmCutlassNvfp4Impl,
@@ -62,7 +66,7 @@ def cutlass_leaves_in_resolution_order(
 ) -> tuple[type[CutlassFusedMoEBase], ...]:
     """Every leaf publishing ``quant_algo``, in ``CUTLASS_LEAVES`` order.
 
-    One entry for most formats and two for ``fp8_block_scales``, so a caller
+    One entry for most formats and three for ``fp8_block_scales``, so a caller
     asking "would a run serve this?" has to ask each, as resolution does.
     """
     quant = normalize_quant(canonical_quant(quant_algo))
@@ -75,9 +79,11 @@ def find_cutlass_leaf(
     """The leaf implementing ``quant_algo``, or ``None`` if none publishes it.
 
     A lookup, not an eligibility check, and resolution does not use it. ``sm``
-    only picks between the two ``fp8_block_scales`` leaves and defaults to the
-    SM resolution sees (``collect_moe_environment``, so an
-    ``override_moe_environment`` applies here too).
+    only picks among the ``fp8_block_scales`` leaves and defaults to the SM
+    resolution sees (``collect_moe_environment``, so an
+    ``override_moe_environment`` applies here too). A leaf with an
+    ``opt_in_flag`` is returned only when the collected environment sets that
+    flag to ``1``, as resolution would admit it.
 
     Args:
         quant_algo: The layer's quantization; ``None`` means unquantized.
@@ -90,9 +96,18 @@ def find_cutlass_leaf(
     leaves = cutlass_leaves_in_resolution_order(quant_algo)
     if len(leaves) <= 1:
         return leaves[0] if leaves else None
+    environment = collect_moe_environment()
     if sm is None:
-        sm = collect_moe_environment().sm
-    return next((leaf for leaf in leaves if leaf.sm_support.accepts(sm)), None)
+        sm = environment.sm
+    return next(
+        (
+            leaf
+            for leaf in leaves
+            if leaf.sm_support.accepts(sm)
+            and (leaf.opt_in_flag is None or environment.env_flag(leaf.opt_in_flag.value) == "1")
+        ),
+        None,
+    )
 
 
 def cutlass_leaf(
@@ -128,6 +143,7 @@ def cutlass_leaf(
     "TrtllmCutlassW4a16Mxfp4Impl",
     "TrtllmCutlassW4a16Nvfp4Impl",
     "TrtllmCutlassW8a16Impl",
+    "TrtllmTritonFp8ActGroupedGemmFp8BlockScalesImpl",
     "TrtllmTritonFp8BlockScalesImpl",
     "cutlass_leaf",
     "cutlass_leaves_in_resolution_order",

tensorrt_llm/_torch/moe/fused_moe/cutlass/base.py

@@ -37,7 +37,7 @@
 )
 from ..impl_base import MoEImplBase, apply_moe_impl_construction_state
 from ..impl_contract import MoERunContext, require_comm_plan
-from ..impl_environment import collect_moe_environment
+from ..impl_environment import MoEEnvFlag, collect_moe_environment
 from ..quantization import MoEWeightLoadingMode
 from ..routing import BaseMoeRoutingMethod
 from .eligibility import SmSupport
@@ -96,6 +96,7 @@ class CutlassFusedMoEBase(MoEImplBase):
       weight method can load a gpt-oss / MiniMax SwiGLU shape.
     * ``requires_oss_cutlass_build`` / ``uses_dynamic_fc2_scale`` -- whether its
       format, or its dynamic FC2 scale, needs the open-source CUTLASS kernels.
+    * ``opt_in_flag`` -- the environment flag that admits an opt-in leaf.
     * ``can_implement`` and ``_get_quant_method``.
     """
 
@@ -112,6 +113,10 @@ class CutlassFusedMoEBase(MoEImplBase):
     #: Whether ``run_moe`` issues ``torch.ops.trtllm.fused_moe``; only such a
     #: leaf is bound by the open-source CUTLASS build (its gate and all-to-all).
     issues_fused_moe_op: bool = True
+    #: The flag a deployment sets to ``1`` to admit this leaf, or None for a leaf
+    #: that needs none. Its ``can_implement`` checks it, and ``find_cutlass_leaf``
+    #: returns the leaf only when it is set.
+    opt_in_flag: MoEEnvFlag | None = None
 
     FUSED_MOE_OP_FLAGS: FusedMoeOpFlags = FusedMoeOpFlags()
     #: Round ``hidden_size`` and ``intermediate_size_per_partition`` up to a

tensorrt_llm/_torch/moe/fused_moe/cutlass/fp8_block_scales_hopper_triton.py

@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""``trtllm.triton.fp8_act_grouped_gemm.fp8_block_scales``."""
+
+import torch
+
+from ....utils import ActivationType
+from ..activation import ActivationParamShape, MoEActivationSupport
+from ..fused_moe_triton_fp8_block_scale import run_triton_fp8_block_scale_moe_fp8_act
+from ..impl_contract import (
+    MoEDeployment,
+    MoEEligibility,
+    MoEProblem,
+    MoERejectReason,
+    MoERunContext,
+    check_activation_support,
+    require_comm_plan,
+)
+from ..impl_environment import MoEEnvFlag
+from ..impl_identity import register_moe_impl
+from .eligibility import check_cutlass_leaf
+from .fp8_block_scales_hopper import DeepgemmCudaHopperFp8BlockScalesImpl
+from .identity import KERNEL_FP8_ACT_GROUPED_GEMM, TECHNIQUE_TRITON, cutlass_descriptor
+
+
+def _check_opted_in(cls: type, d: MoEDeployment) -> MoEEligibility | None:
+    """Decline unless the deployment opted in through ``cls.opt_in_flag``."""
+    if d.env.env_flag(cls.opt_in_flag.value) == "1":
+        return None
+    return MoEEligibility.no(
+        MoERejectReason.PATH_NOT_ENABLED,
+        f"{cls.__name__} is opt-in: set {cls.opt_in_flag.value}=1 to enable it",
+    )
+
+
+def _check_no_expert_bias(cls: type, p: MoEProblem) -> MoEEligibility | None:
+    """Decline an expert FC bias, which the Triton kernels take no operand for."""
+    if p.bias:
+        return MoEEligibility.no(
+            MoERejectReason.ACTIVATION_UNSUPPORTED, f"{cls.__name__} takes no expert bias"
+        )
+    return None
+
+
+@register_moe_impl
+class TrtllmTritonFp8ActGroupedGemmFp8BlockScalesImpl(DeepgemmCudaHopperFp8BlockScalesImpl):
+    """``trtllm.triton.fp8_act_grouped_gemm.fp8_block_scales``.
+
+    FP8 block scales on SM90 for decode-sized batches. The grouped DeepGEMM
+    kernel of the parent leaf walks every expert group on one persistent CTA
+    per SM; the Triton kernels here quantize the activations per 1 x 128 group
+    and launch only as many expert blocks as the batch's (token, expert) pairs
+    can occupy. A call takes the Triton path when all of these hold
+    (:meth:`_takes_triton_path`); otherwise it runs the parent's fused op:
+
+    - ``run_moe`` receives at most ``TRITON_FP8_ACT_MAX_TOKENS`` rows. These
+      are the rows after dispatch, so with attention DP and an allgather they
+      count every DP rank's tokens.
+    - The scheduler hands no ``moe_output`` buffer. The Triton path returns its
+      own tensor, so it never runs under a communication strategy that
+      combines from a workspace (NVLinkOneSided).
+    - BF16 input without scale factors.
+    - The layer's clamp, if any, is one value for every local expert
+      (decided once in :meth:`cache_derived_state`).
+
+    Opt-in: resolution passes this leaf over unless
+    ``TLLM_MOE_ENABLE_SM90_TRITON_FP8_DECODE=1`` (``opt_in_flag``), so by
+    default the parent leaf serves every SM90 FP8 block-scale layer.
+
+    With expert parallelism, (token, expert) pairs owned by other ranks keep a
+    clamped local expert id with a zero weight, so they are still computed.
+    """
+
+    #: Largest number of rows ``run_moe`` receives that runs the Triton path.
+    TRITON_FP8_ACT_MAX_TOKENS = 32
+    #: The environment flag that admits this leaf; read through ``d.env``.
+    opt_in_flag = MoEEnvFlag.SM90_TRITON_FP8_DECODE
+
+    descriptor = cutlass_descriptor(
+        "fp8_block_scales",
+        "Triton FP8-activation grouped GEMM for decode-sized batches, the DeepGEMM grouped "
+        "GEMM otherwise, SM90.",
+        technique=TECHNIQUE_TRITON,
+        kernel_name=KERNEL_FP8_ACT_GROUPED_GEMM,
+        # The Triton kernels take no bias operands.
+        supports_expert_bias=False,
+    )
+    # The Triton kernels run SwiGLU only. The clamp stays a per-expert tensor
+    # for the fused op; the Triton path takes it only when it is uniform.
+    activation_support = MoEActivationSupport(
+        kinds=frozenset({ActivationType.Swiglu}),
+        limit=ActivationParamShape.PER_EXPERT_TENSOR,
+        clamp_after_silu=True,
+    )
+
+    #: ``(limit, clamp_after_silu)`` for the Triton SwiGLU, or None while the
+    #: layer cannot take the Triton path. Set by :meth:`cache_derived_state`.
+    _triton_swiglu: tuple[float | None, bool] | None = None
+
+    @classmethod
+    def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility:
+        return check_cutlass_leaf(
+            cls,
+            p,
+            d,
+            check_activation_support(cls, p),
+            _check_no_expert_bias(cls, p),
+            _check_opted_in(cls, d),
+        )
+
+    def cache_derived_state(self) -> None:
+        super().cache_derived_state()
+        self._triton_swiglu = self._uniform_swiglu_clamp()
+
+    def _uniform_swiglu_clamp(self) -> tuple[float | None, bool] | None:
+        """The Triton SwiGLU settings read off the installed ``act_*`` slots.
+
+        Returns ``(limit, clamp_after_silu)``, or None when the layer has an
+        expert bias, alpha/beta constants, or a clamp that differs between
+        experts: the Triton kernel takes the clamp by value.
+        """
+        if (
+            self.activation_type != ActivationType.Swiglu
+            or self.w3_w1_bias is not None
+            or self.w2_bias is not None
+            or self.act_alpha is not None
+            or self.act_beta is not None
+        ):
+            return None
+        clamp = self.act_clamp
+        if clamp is None:
+            return None, False
+        if isinstance(clamp, torch.Tensor):
+            flat = clamp.detach().reshape(-1)
+            if flat.numel() == 0 or not bool((flat == flat[0]).all()):
+                return None
+            clamp = flat[0].item()
+        return float(clamp), bool(self.act_clamp_after_silu)
+
+    def _takes_triton_path(self, ctx: MoERunContext) -> bool:
+        """Whether this call runs the Triton kernels; see the class docstring."""
+        return (
+            self._triton_swiglu is not None
+            and ctx.x.shape[0] <= self.TRITON_FP8_ACT_MAX_TOKENS
+            and ctx.x.dtype == torch.bfloat16
+            and ctx.x_sf is None
+            and require_comm_plan(self, ctx).moe_output is None
+        )
+
+    def run_moe(self, ctx: MoERunContext, *, workspace: dict | None = None) -> torch.Tensor:
+        if not self._takes_triton_path(ctx):
+            return super().run_moe(ctx, workspace=workspace)
+        del workspace  # The Triton path allocates its own intermediates.
+
+        # The scheduler leaves token_final_scales as None when the router
+        # weights were already folded into x.
+        token_final_scales = ctx.token_final_scales
+        if token_final_scales is None:
+            token_final_scales = torch.ones_like(ctx.token_selected_experts, dtype=torch.float32)
+        if self.ep_size == 1:
+            # Every routed slot is local: the IDs and weights need no remap.
+            local_ids, local_scales = ctx.token_selected_experts, token_final_scales
+        else:
+            # The kernel indexes local experts; zero the scale of every pair
+            # this rank does not own.
+            local_ids, offsets = self._local_slot_ids(ctx.token_selected_experts)
+            local_scales = token_final_scales * (local_ids == offsets).to(token_final_scales.dtype)
+
+        swiglu_limit, swiglu_clamp_after_silu = self._triton_swiglu
+        return run_triton_fp8_block_scale_moe_fp8_act(
+            ctx.x,
+            local_ids,
+            local_scales,
+            self.w3_w1_weight,
+            self.quant_scales.fc_weight_scales,
+            self.w2_weight,
+            self.quant_scales.proj_weight_scales,
+            swiglu_limit=swiglu_limit,
+            swiglu_clamp_after_silu=swiglu_clamp_after_silu,
+            output_dtype=ctx.output_dtype,
+        )

tensorrt_llm/_torch/moe/fused_moe/cutlass/identity.py

@@ -16,7 +16,8 @@
 
 Ten leaves are ``trtllm.cutlass.grouped_gemm.<quant>``. Each ``fp8_block_scales``
 leaf names the kernel it runs instead: ``deepgemm.cuda.hopper_grouped_gemm``
-(:mod:`.fp8_block_scales_hopper`) and ``trtllm.triton.grouped_gemm``
+(:mod:`.fp8_block_scales_hopper`), ``trtllm.triton.fp8_act_grouped_gemm``
+(:mod:`.fp8_block_scales_hopper_triton`) and ``trtllm.triton.grouped_gemm``
 (:mod:`.fp8_block_scales_triton`).
 """
 
@@ -37,6 +38,8 @@
 
 KERNEL_GROUPED_GEMM = "grouped_gemm"
 KERNEL_HOPPER_GROUPED_GEMM = "hopper_grouped_gemm"
+#: Grouped GEMM over FP8 activations quantized per 1 x 128 group.
+KERNEL_FP8_ACT_GROUPED_GEMM = "fp8_act_grouped_gemm"
 
 CUTLASS_CAPABILITIES = MoEStaticCapability(
     supports_expert_bias=True,
@@ -63,8 +66,8 @@ def cutlass_descriptor(
         quant: The one quantization format the leaf implements.
         doc: One line naming the kernel.
         provider: Identity segment. The defaults name the CUTLASS grouped
-            GEMM; only the two FP8 block-scale leaves, which run another
-            kernel, override the segments.
+            GEMM; only the FP8 block-scale leaves, which run other kernels,
+            override the segments.
         technique: Identity segment; see ``provider``.
         kernel_name: Identity segment; see ``provider``.
         supports_moe_lora: Whether ``torch.ops.trtllm.fused_moe`` fuses