diff --git a/py/torch_tensorrt/dynamo/conversion/impl/linear.py b/py/torch_tensorrt/dynamo/conversion/impl/linear.py index 3827284950..4adc75e343 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/linear.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/linear.py @@ -8,7 +8,7 @@ from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext -from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor +from torch_tensorrt.dynamo.conversion.converter_utils import broadcast, get_trt_tensor def linear( @@ -49,6 +49,18 @@ def linear( ) if bias is not None: + # Rank-align a 1-D linear bias with the matmul output while preserving + # singleton dimensions. TensorRT can broadcast those dimensions in the + # elementwise add without materializing the bias to the full output + # shape (for example, [O] -> [1, 1, O] for a [B, S, O] output). + out, bias = broadcast( + ctx, + out, + bias, + f"{name}_add_bias_lhs", + f"{name}_add_bias_rhs", + ) + # add bias out = impl.elementwise.add( ctx, target, source_ir, f"{name}_add_bias", out, bias diff --git a/tests/py/dynamo/conversion/test_linear_aten.py b/tests/py/dynamo/conversion/test_linear_aten.py index 2426b7b42d..8619bebd7d 100644 --- a/tests/py/dynamo/conversion/test_linear_aten.py +++ b/tests/py/dynamo/conversion/test_linear_aten.py @@ -49,6 +49,19 @@ def forward(self, x, weight, bias): LinearModel(), input_specs, use_dynamo_tracer=True, enable_passes=True ) + def test_linear_with_rank_3_input_and_bias(self): + class LinearModel(torch.nn.Module): + def forward(self, x, weight, bias): + return torch.ops.aten.linear.default(x, weight, bias) + + inputs = [ + torch.randn(2, 3, 4).cuda(), + torch.randn(5, 4).cuda(), + torch.randn(5).cuda(), + ] + + self.run_test(LinearModel(), inputs, use_dynamo_tracer=True, enable_passes=True) + if __name__ == "__main__": run_tests()