Fix linear bias broadcasting in the Dynamo converter

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

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()