[None][perf] Fuse WideEP post-MoE add and RMSNorm

tensorrt_llm/_torch/compilation/utils.py

@@ -65,6 +65,13 @@ def capture_piecewise_cuda_graph(enable: bool):
 
 def inplace_info():
     inplace_map = {
+        # These keys index auto_functionalized's mutated-output tuple, not the
+        # raw operator arguments. ``additional`` is read-only, so the mutated
+        # ``input`` and ``residual`` tensors occupy output slots 1 and 2.
+        torch.ops.trtllm.flashinfer_fused_add_add_rmsnorm.default: {
+            1: "input",
+            2: "residual"
+        },
         torch.ops.trtllm.flashinfer_fused_add_rmsnorm.default: {
             1: "input",
             2: "residual"

tensorrt_llm/_torch/custom_ops/__init__.py

@@ -48,13 +48,14 @@ def inplace_slice_copy(dest: torch.Tensor, src: torch.Tensor, dim1_start: int,
 if IS_FLASHINFER_AVAILABLE:
     from .flashinfer_custom_ops import (
         flashinfer_apply_rope_with_cos_sin_cache_inplace,
-        flashinfer_fused_add_rmsnorm, flashinfer_gelu_tanh_and_mul,
-        flashinfer_gemma_fused_add_rmsnorm, flashinfer_gemma_rmsnorm,
-        flashinfer_rmsnorm, flashinfer_silu_and_mul)
+        flashinfer_fused_add_add_rmsnorm, flashinfer_fused_add_rmsnorm,
+        flashinfer_gelu_tanh_and_mul, flashinfer_gemma_fused_add_rmsnorm,
+        flashinfer_gemma_rmsnorm, flashinfer_rmsnorm, flashinfer_silu_and_mul)
     __all__ += [
         'flashinfer_gelu_tanh_and_mul',
         'flashinfer_silu_and_mul',
         'flashinfer_rmsnorm',
+        'flashinfer_fused_add_add_rmsnorm',
         'flashinfer_fused_add_rmsnorm',
         'flashinfer_apply_rope_with_cos_sin_cache_inplace',
         'flashinfer_gemma_fused_add_rmsnorm',

tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py

@@ -63,6 +63,50 @@ def flashinfer_fused_add_rmsnorm(input: torch.Tensor,
                           eps,
                           enable_pdl=get_env_enable_pdl())
 
+    @torch.library.custom_op("trtllm::flashinfer_fused_add_add_rmsnorm",
+                             mutates_args=("input", "residual"))
+    def flashinfer_fused_add_add_rmsnorm(
+            input: torch.Tensor,  # noqa: A002 - part of the public op schema
+            additional: torch.Tensor,
+            residual: torch.Tensor,
+            weight: torch.Tensor,
+            eps: float) -> None:
+        """Fuse a BF16 MoE add, residual add, and RMSNorm in place.
+
+        Args:
+            input: Contiguous or row-strided BF16 tensor of shape ``(M, H)``.
+                Updated in place with the normalized output.
+            additional: BF16 tensor of shape ``(M, H)`` containing the second
+                MoE contribution. This tensor is not modified.
+            residual: BF16 tensor of shape ``(M, H)``. Updated in place with
+                the BF16-rounded MoE sum plus the original residual.
+            weight: Contiguous BF16 RMSNorm weight of shape ``(H,)``.
+            eps: Epsilon added to the RMSNorm variance.
+
+        Returns:
+            ``None``. ``input`` and ``residual`` are updated in place.
+        """
+        # Keep this import lazy so installations that use FlashInfer's CUDA
+        # norm fallback do not need CUTLASS DSL unless the default-off WideEP
+        # path is explicitly selected.
+        from ..cute_dsl_kernels.flashinfer_fused_add_add_rmsnorm import \
+            fused_add_add_rmsnorm_cute
+        fused_add_add_rmsnorm_cute(input,
+                                   additional,
+                                   residual,
+                                   weight,
+                                   eps,
+                                   enable_pdl=get_env_enable_pdl())
+
+    @flashinfer_fused_add_add_rmsnorm.register_fake
+    def _(
+            input: torch.Tensor,  # noqa: A002 - mirrors the public op schema
+            additional: torch.Tensor,
+            residual: torch.Tensor,
+            weight: torch.Tensor,
+            eps: float) -> None:
+        pass
+
     @torch.library.custom_op("trtllm::flashinfer_fused_add_rmsnorm_quant",
                              mutates_args=("out", "residual"))
     def flashinfer_fused_add_rmsnorm_quant(out: torch.Tensor,