[TRTLLM-15398][perf] VisualGen MLP: cublasLt GELU-tanh epilogue for the unquantized bf16 path

tensorrt_llm/_torch/modules/mlp.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from collections.abc import Callable
 from typing import Optional, Tuple, Union
 
@@ -10,8 +13,9 @@
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType, add_lora_result
 from ..utils import Fp4QuantizedTensor, gelu_tanh, relu2
-from .linear import (Linear, TensorParallelMode, WeightMode,
-                     WeightsLoadingConfig, is_static_nvfp4_input_eligible)
+from .linear import (Linear, TensorParallelMode, UnquantizedLinearMethod,
+                     WeightMode, WeightsLoadingConfig,
+                     is_static_nvfp4_input_eligible)
 
 
 class MLP(nn.Module):
@@ -115,7 +119,7 @@ def create_weights(self):
         # Static eligibility for the fused GELU(tanh) CuteDSL epilogue (mirrors
         # GatedMLP); the runtime quant_method check is deferred to first forward.
         self._use_fused_gelu, self._use_fused_gelu_fp4out = (
-            self._gelu_fusion_eligibility())
+            self._nvfp4_gelu_fusion_eligibility())
 
     # Minimum M for the fp4out CuTe DSL GELU kernel; below this its SFC epilogue
     # can write out-of-bounds (CTA tile height > output rows), so fall back to
@@ -148,29 +152,61 @@ def forward(
                     self._fused_gelu(x, fp4_out=m >= MLP._FP4OUT_MIN_M))
             return self.down_proj(self._fused_gelu(x))
 
-        x_up = self.up_proj(x)
-
-        # Weight loading may replace the quantization method after
-        # create_weights(), so do not rely on the cached eligibility alone.
-        if (self._use_fused_relu2_quant
-                and is_static_nvfp4_input_eligible(self.down_proj)):
-            x_act = self._fused_relu2_quant(x_up)
+        if self._unquantized_gelu_fusion_eligible(x):
+            x_act = self._fused_up_proj_gelu(x)
         else:
-            x_act = self.activation(x_up)
+            x_up = self.up_proj(x)
+            # Weight loading may replace the quantization method after
+            # create_weights(), so do not rely on the cached eligibility alone.
+            if (self._use_fused_relu2_quant
+                    and is_static_nvfp4_input_eligible(self.down_proj)):
+                x_act = self._fused_relu2_quant(x_up)
+            else:
+                x_act = self.activation(x_up)
 
         x_down = self.down_proj(x_act)
 
         return x_down
 
-    def _gelu_fusion_eligibility(self) -> Tuple[bool, bool]:
+    def _unquantized_gelu_fusion_eligible(self, x: torch.Tensor) -> bool:
+        """Whether the unquantized up projection can use the cuBLASLt GELU
+        epilogue without bypassing Linear post-processing or another GEMM
+        backend.
+        """
+        up_proj = self.up_proj
+        return (self.activation is gelu_tanh
+                and hasattr(torch, "_addmm_activation")
+                and isinstance(x, torch.Tensor) and x.dim() >= 2 and x.is_cuda
+                and x.dtype == torch.bfloat16 and not torch.is_grad_enabled()
+                and type(up_proj.quant_method) is UnquantizedLinearMethod
+                and up_proj.bias is not None
+                and up_proj.weight.dtype == torch.bfloat16
+                and not up_proj.gather_output
+                and not up_proj.use_custom_cublas_mm
+                and not up_proj.use_cute_dsl_bf16_gemm)
+
+    def _fused_up_proj_gelu(self, x: torch.Tensor) -> torch.Tensor:
+        """Run the bf16 up projection with a fused GELU(tanh) epilogue."""
+        input_shape = x.shape
+        # VisualGen diffusion transformers commonly process long token
+        # sequences. Folding GELU into the up-projection avoids an extra
+        # full-tensor read/write and standalone activation kernel at large M.
+        output = torch._addmm_activation(self.up_proj.bias,
+                                         x.reshape(-1, input_shape[-1]),
+                                         self.up_proj.weight.t(),
+                                         use_gelu=True)
+        return output.reshape(*input_shape[:-1], output.shape[-1])
+
+    def _nvfp4_gelu_fusion_eligibility(self) -> Tuple[bool, bool]:
         """Return (bf16_out_ok, fp4_out_ok) static eligibility for the fused
         GELU(tanh) epilogue (mirrors GatedMLP's SwiGLU paths). Requires the
-        Blackwell CuteDSL op(s), SM 100/103, and an NVFP4 up_proj; fp4-out builds
-        on bf16-out and also needs an NVFP4 down_proj with a static input_scale
-        and no forced dynamic quantization. The runtime quant_method check is
-        applied in forward (quant_method can be downgraded after this).
+        Blackwell CuteDSL op(s), SM 100/103, a local (not gathered) NVFP4
+        up_proj output; fp4-out builds on bf16-out and also needs an NVFP4
+        down_proj with a static input_scale and no forced dynamic quantization.
+        The runtime quant_method check is applied in forward (quant_method can
+        be downgraded after this).
         """
-        if (self.activation is not gelu_tanh
+        if (self.activation is not gelu_tanh or self.up_proj.gather_output
                 or get_sm_version() not in (100, 103) or not getattr(
                     self.up_proj, "has_nvfp4_activation_quantization", False)):
             return False, False

tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py

@@ -21,6 +21,8 @@
 fall back to below that floor.
 """
 
+from unittest import mock
+
 import pytest
 import torch
 import torch.nn.functional as F
@@ -35,6 +37,165 @@
 SF_VEC = 16
 
 
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not hasattr(torch, "_addmm_activation"),
+    reason="requires CUDA torch._addmm_activation",
+)
+def test_mlp_gelu_tanh_backward_uses_unfused_path() -> None:
+    """Autograd keeps eager linear + GELU so bf16 backward remains valid."""
+    from tensorrt_llm._torch.modules.mlp import MLP
+    from tensorrt_llm._torch.utils import gelu_tanh
+
+    mlp = MLP(
+        hidden_size=64,
+        intermediate_size=128,
+        bias=True,
+        activation=gelu_tanh,
+        dtype=torch.bfloat16,
+        reduce_output=False,
+    ).cuda()
+    with torch.no_grad():
+        for parameter in mlp.parameters():
+            parameter.normal_(std=0.02)
+
+    x = torch.randn(2, 4, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+    with mock.patch.object(
+        torch,
+        "_addmm_activation",
+        side_effect=AssertionError("inference-only fusion used with autograd"),
+    ):
+        mlp(x).float().sum().backward()
+
+    assert x.grad is not None
+    assert torch.isfinite(x.grad).all()
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not hasattr(torch, "_addmm_activation"),
+    reason="requires CUDA torch._addmm_activation",
+)
+@pytest.mark.parametrize(
+    "excluded_linear_path",
+    ["gather_output", "use_custom_cublas_mm", "use_cute_dsl_bf16_gemm"],
+)
+def test_mlp_gelu_tanh_excluded_linear_path_uses_unfused_path(
+    excluded_linear_path: str,
+) -> None:
+    """Linear post-processing and alternate GEMM backends keep their dispatch."""
+    from tensorrt_llm._torch.modules.mlp import MLP
+    from tensorrt_llm._torch.utils import gelu_tanh
+
+    mlp = MLP(
+        hidden_size=64,
+        intermediate_size=128,
+        bias=True,
+        activation=gelu_tanh,
+        dtype=torch.bfloat16,
+        reduce_output=False,
+    )
+    setattr(mlp.up_proj, excluded_linear_path, True)
+
+    x = torch.randn(2, 4, 64, dtype=torch.bfloat16, device="cuda")
+    x_up = torch.randn(2, 4, 128, dtype=torch.bfloat16, device="cuda")
+    x_down = torch.randn_like(x)
+    with (
+        torch.no_grad(),
+        mock.patch.object(
+            mlp,
+            "_fused_up_proj_gelu",
+            side_effect=AssertionError("fusion bypassed Linear dispatch"),
+        ),
+        mock.patch.object(mlp.up_proj, "forward", return_value=x_up) as up_forward,
+        mock.patch.object(mlp.down_proj, "forward", return_value=x_down),
+    ):
+        output = mlp(x)
+
+    up_forward.assert_called_once_with(x)
+    assert output is x_down
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or not hasattr(torch, "_addmm_activation"),
+    reason="requires CUDA torch._addmm_activation",
+)
+def test_mlp_gelu_tanh_eligible_path_uses_fused_epilogue() -> None:
+    """An eligible bf16 inference MLP engages the fused epilogue and stays
+    close to the unfused reference."""
+    from tensorrt_llm._torch.modules.mlp import MLP
+    from tensorrt_llm._torch.utils import gelu_tanh
+
+    mlp = MLP(
+        hidden_size=64,
+        intermediate_size=128,
+        bias=True,
+        activation=gelu_tanh,
+        dtype=torch.bfloat16,
+        reduce_output=False,
+    ).cuda()
+    with torch.no_grad():
+        for parameter in mlp.parameters():
+            parameter.normal_(std=0.02)
+
+    x = torch.randn(2, 4, 64, dtype=torch.bfloat16, device="cuda")
+    with (
+        torch.no_grad(),
+        mock.patch.object(mlp, "_fused_up_proj_gelu", wraps=mlp._fused_up_proj_gelu) as fused,
+    ):
+        fused_out = mlp(x)
+    fused.assert_called_once_with(x)
+
+    with (
+        torch.no_grad(),
+        mock.patch.object(mlp, "_unquantized_gelu_fusion_eligible", return_value=False),
+    ):
+        unfused_out = mlp(x)
+
+    assert fused_out.shape == unfused_out.shape
+    torch.testing.assert_close(fused_out.float(), unfused_out.float(), atol=2e-2, rtol=2e-2)
+
+
+def test_mlp_nvfp4_gelu_gather_output_is_ineligible() -> None:
+    """The direct NVFP4 GELU kernels must not bypass column all-gather."""
+    from tensorrt_llm._torch.modules.mlp import MLP
+    from tensorrt_llm._torch.utils import gelu_tanh
+
+    mlp = MLP(
+        hidden_size=64,
+        intermediate_size=128,
+        bias=True,
+        activation=gelu_tanh,
+        dtype=torch.bfloat16,
+        reduce_output=False,
+    )
+    up = torch.nn.Identity()
+    up.gather_output = True
+    up.has_nvfp4_activation_quantization = True
+    mlp.up_proj = up
+    down = torch.nn.Identity()
+    down.has_nvfp4_activation_quantization = True
+    down.force_dynamic_quantization = False
+    down.input_scale = torch.ones(1)
+    down.pre_quant_scale = None
+    mlp.down_proj = down
+
+    with (
+        mock.patch("tensorrt_llm._torch.modules.mlp.get_sm_version", return_value=100),
+        mock.patch.object(
+            torch.ops.trtllm,
+            "cute_dsl_nvfp4_dense_gemm_gelu_blackwell",
+            object(),
+            create=True,
+        ),
+        mock.patch.object(
+            torch.ops.trtllm,
+            "cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell",
+            object(),
+            create=True,
+        ),
+    ):
+        assert mlp._nvfp4_gelu_fusion_eligibility() == (False, False)
+
+
 def _quantize_nvfp4(x_bf16: torch.Tensor):
     """Quantize [., K] bf16 -> (fp4 packed, swizzled SF, global_sf scalar)."""
     global_sf = x_bf16.abs().max().float() / (FP8_E4M3_MAX * FP4_E2M1_MAX)

tests/unittest/scripts/test_cbts_coverage_pilot.py

@@ -146,7 +146,7 @@ def urlopen(_request: urllib.request.Request, *, timeout: int) -> NoReturn:
 def test_main_reads_bot_trigger_payload(
     pilot_module: ModuleType,
     monkeypatch: pytest.MonkeyPatch,
-    capsys: pytest.CaptureFixture[str],
+    capfd: pytest.CaptureFixture[str],
 ) -> None:
     trigger_phrase = json.dumps({"github_pr_api_url": PR_API_URL})
     monkeypatch.setenv("gitlabTriggerPhrase", trigger_phrase)
@@ -164,6 +164,6 @@ def check_pilot_eligibility(
     monkeypatch.setattr(pilot_module, "check_pilot_eligibility", check_pilot_eligibility)
 
     assert pilot_module.main([]) == 0
-    captured = capsys.readouterr()
+    captured = capfd.readouterr()
     assert captured.out == "true\n"
     assert "pr_author=pilot-user, eligible=true" in captured.err