diff --git a/tensorrt_llm/_torch/modules/mlp.py b/tensorrt_llm/_torch/modules/mlp.py index bad1078027b6..a13476bb5137 100644 --- a/tensorrt_llm/_torch/modules/mlp.py +++ b/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 diff --git a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py index 586973673f87..4fb55ed93d76 100644 --- a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py +++ b/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) diff --git a/tests/unittest/scripts/test_cbts_coverage_pilot.py b/tests/unittest/scripts/test_cbts_coverage_pilot.py index 845009fbccd0..1ec7266b97ad 100644 --- a/tests/unittest/scripts/test_cbts_coverage_pilot.py +++ b/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