Fuse pad into convolution

[micwill755]

Description

Fold asymmetric constant_pad_nd → convolution patterns into TensorRT convolution pre_padding / post_padding via a new Dynamo post-lowering pass (fuse_pad_into_convolution) and tensorrt::conv_asym_pad converter.

Motivation: Wan-style causal 3D convolutions (used by Cosmos3 / Cosmos3-Edge VAE) zero their own padding and apply it themselves, often asymmetrically in time (e.g. (2 * pad_t, 0)). TensorRT’s ONNX parser already folds Pad → Conv into the convolution’s explicit pre/post padding. Torch-TRT previously only set symmetric padding_nd, so those pads stayed as materialized copy kernels (extra work on the VAE encode path).

Change:

Post-lowering pass rewrites eligible constant_pad_nd → convolution into tensorrt::conv_asym_pad
Converter emits IConvolutionLayer with explicit pre_padding / post_padding (via extended convNd)
Unit tests in tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py
Validation on Cosmos3-Edge policy (Torch-TRT → cosmos3_policy_inference):

Geometry: 256×256, 5 frames, und_len 121, batch 1, 4 denoise steps
Pad fusion removed ~45 extra VAE pad kernels and improved VAE encode 7.46 → 5.50 ms (~1.96 ms)
Full policy steady-state median improved ~23.85 → 21.81 ms (~2.0 ms) after enabling this pass (with MoT builder settings already applied)
No new runtime dependencies beyond existing Torch-TensorRT / TensorRT.

Fixes # (n/a — performance / parity with ONNX Pad→Conv folding)

Type of change
New feature (non-breaking change which adds functionality)
Performance optimization for models with asymmetric Pad→Conv (e.g. Cosmos3 VAE)
Checklist:
My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 02:49:49.445228+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 02:50:08.060518+00:00
@@ -206,11 +206,13 @@
        x = torch.randn(1, 4, 5, 8, 8)
        weight = torch.randn(4, 4, 3, 3, 3)
        bias = torch.randn(4)
        pre, post = [2, 1, 1], [0, 1, 1]
        expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
-        got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+        got = tensorrt_conv_asym_pad_op(
+            x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+        )
        torch.testing.assert_close(got, expected)


class TestFusePadIntoConvolutionConverter(DispatchTestCase):
    """End-to-end TRT numerical tests exercising the pad→conv fusion."""

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 02:49:49.054179+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 02:50:15.101134+00:00
@@ -206,11 +206,13 @@
        x = torch.randn(1, 4, 5, 8, 8)
        weight = torch.randn(4, 4, 3, 3, 3)
        bias = torch.randn(4)
        pre, post = [2, 1, 1], [0, 1, 1]
        expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
-        got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+        got = tensorrt_conv_asym_pad_op(
+            x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+        )
        torch.testing.assert_close(got, expected)


class TestFusePadIntoConvolutionConverter(DispatchTestCase):
    """End-to-end TRT numerical tests exercising the pad→conv fusion."""

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 02:50:10.421070+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 02:50:40.458077+00:00
@@ -206,11 +206,13 @@
        x = torch.randn(1, 4, 5, 8, 8)
        weight = torch.randn(4, 4, 3, 3, 3)
        bias = torch.randn(4)
        pre, post = [2, 1, 1], [0, 1, 1]
        expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
-        got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+        got = tensorrt_conv_asym_pad_op(
+            x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+        )
        torch.testing.assert_close(got, expected)


class TestFusePadIntoConvolutionConverter(DispatchTestCase):
    """End-to-end TRT numerical tests exercising the pad→conv fusion."""

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 03:56:45.218617+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py	2026-07-29 03:57:10.275509+00:00
@@ -206,11 +206,13 @@
        x = torch.randn(1, 4, 5, 8, 8)
        weight = torch.randn(4, 4, 3, 3, 3)
        bias = torch.randn(4)
        pre, post = [2, 1, 1], [0, 1, 1]
        expected = F.conv3d(F.pad(x, [1, 1, 1, 1, 2, 0]), weight, bias, 1, 0, 1, 1)
-        got = tensorrt_conv_asym_pad_op(x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1)
+        got = tensorrt_conv_asym_pad_op(
+            x, weight, bias, [1, 1, 1], pre, post, [1, 1, 1], 1
+        )
        torch.testing.assert_close(got, expected)


class TestFusePadIntoConvolutionConverter(DispatchTestCase):
    """End-to-end TRT numerical tests exercising the pad→conv fusion."""

[narendasan]

@micwill755 run precommit for the linting

[zewenli98]

LGTM overall. just some minor issues

[zewenli98]

Consider moving the validation to the function entry to fail faster for reducing build time.

[zewenli98]

what if pre_padding is None? len(pre_padding) seems problematic. Can we directly pass in an empty list given that pre_padding and post_padding are provided?