From 2e84e298e7b888d20a6c60db50bcc7012bd681cd Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Wed, 26 Aug 2026 06:57:03 -0700 Subject: [PATCH] fix: guard the 3D convolutions the pad-folding pass hides from the conv guard convolution_capability_validator rejects forward 3D convolutions when Turing (SM 7.5) is a build target, because TensorRT-RTX builds an engine for them whose createExecutionContext() then returns nullptr. It is registered on aten.convolution.default only. The fuse_pad_into_convolution lowering pass rewrites constant_pad_nd -> convolution into tensorrt::conv_asym_pad, and that op's converter had no capability_validator at all. By the time partitioning runs there is no aten.convolution node left to validate, so a padded 3D convolution reaches TensorRT-RTX regardless of the guard and surfaces at runtime as Expected exec_ctx_.get() != nullptr to be true but got false Unable to (re)create TensorRT execution context This is a production gap, not just a test gap. The pass runs in the normal compile path, and it fires for any non-transposed, zero-fill, non-negative constant_pad_nd -> conv pair -- asymmetry is not required -- so any real model with a padded conv3d hits it on Turing. * aten_ops_converters.py: extract turing_rejects_forward_convolution() so the "3D forward conv is unsupported on this target" rule lives in one place. convolution_capability_validator delegates to it; its behaviour is unchanged, including failing open when meta["val"] is absent. * custom_ops_converters.py: add conv_asym_pad_capability_validator and register it on the fused op. It reads the spatial rank off the argument list rather than node.meta -- the fused op's args are (source, weight, bias, stride, pre_padding, post_padding, dilation, groups), so len(stride) is the rank -- which makes the guard hold whether or not the caller ran the dynamo tracer. The pass never fuses a transposed convolution, so there is no deconvolution case to handle here. Only rank 3 is rejected; 2D keeps running on TensorRT. * test_fuse_pad_into_convolution.py: skip test_padded_conv3d on Turing. DispatchTestCase has no PyTorch-fallback path -- run_test hands the graph straight to TRTInterpreter, skipping the partitioner -- so a rejected node raises UnsupportedOperatorException instead of falling back. Guarding without skipping would only swap one failure for another. * test_turing_capability_guards.py: add a padded-conv3d fallback case and a padded-conv2d positive control. Both key on target_compute_capabilities=[(7, 5)] rather than the live device, so they exercise the guard on any GPU and therefore in CI. Reverting just the registration makes the conv3d case fail, so it covers the gap rather than merely restating it. Deliberately not fixed in the lowering pass. Declining to fuse when the underlying convolution would be rejected keeps one rule in one place and would pick up any future guard on aten.convolution for free, and the pass already receives the settings it would need. But TestFusePadIntoConvolutionPass builds its graphs with a default CompilationSettings(), and test_graph_contains_fused_op_after_lowering calls post_lowering with one too. Default settings mean target_compute_capabilities=None, which resolves against the current device -- so on a Turing GPU the pass would decline to fuse and the five tests that assert the fusion *does* happen would fail. Each would then need a pinned non-Turing target or a skip. Testing (T4 / SM 7.5 and L40S / SM 8.9, driver 595.58.03, identical pinned stacks): * All 8 test_padded_conv3d cases go fail -> skip on Turing and still run and pass on the L40S control, confirming the guard is inert off Turing. * Full lowering/ sweep on both arms, all 259 tests accounted for on each. On Turing exactly 8 tests change status, all of them fail -> skip and all of them test_padded_conv3d: 9 -> 1 failed, 4 -> 12 skipped, and 242 passing either way. (The other 4 pre-existing skips, and 4 fp16 SDPA tests that abort the worker, are unrelated and unchanged.) On the non-Turing control every one of the 259 keeps its previous status and the skip count stays at 0. * The 11 other tests in that file stay green on both arms -- in particular test_graph_contains_fused_op_after_lowering and test_fuses_causal_3d_pad, both of which build 3D graphs and assert the fusion still happens on the T4. * test_turing_capability_guards.py passes on both arms: T4 17 -> 19 passed / 1 skipped, L40S 12 -> 14 passed / 6 skipped. Co-Authored-By: Claude Opus 5 --- .../dynamo/conversion/aten_ops_converters.py | 42 +++++++++++++----- .../conversion/custom_ops_converters.py | 38 ++++++++++++++-- .../test_fuse_pad_into_convolution.py | 7 ++- .../models/test_turing_capability_guards.py | 44 +++++++++++++++++++ 4 files changed, 115 insertions(+), 16 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index cc679ec6c2..4b7385438a 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -3252,6 +3252,30 @@ def aten_ops_le( _CONV_ARG_TRANSPOSED = 6 +def turing_rejects_forward_convolution( + node: Node, + spatial_rank: Optional[int], + settings: Optional[CompilationSettings] = None, +) -> bool: + """Whether a forward convolution over ``spatial_rank`` spatial dims must fall back. + + No valid kernel config for 3D ConvFwd on SM 7.5 for TensorRT-RTX, a known gap. + Transposed 3D is a distinct layer and is unaffected, so callers must only ask about + non-transposed convolutions. + """ + unsupported_spatial_rank = 3 + if spatial_rank != unsupported_spatial_rank: + return False + if not trt_rtx_targets_turing(settings): + return False + _LOGGER.debug( + "3D convolution '%s' is not supported on TensorRT-RTX for Turing " + "(SM 7.5). Falling back to PyTorch.", + node.name, + ) + return True + + def convolution_capability_validator( node: Node, settings: Optional[CompilationSettings] = None ) -> bool: @@ -3278,20 +3302,14 @@ def convolution_capability_validator( ) return False - # No valid kernel config for 3D ConvFwd on SM 7.5: the engine builds, but - # createExecutionContext() then returns nullptr. Transposed 3D is a distinct layer - # and is unaffected. aten.convolution input is (N, C, *spatial), so ndim 5 is 3D. - is_forward_conv = not args_bounds_check(node.args, _CONV_ARG_TRANSPOSED) - if trt_rtx_targets_turing(settings) and is_forward_conv: + # aten.convolution input is (N, C, *spatial), so ndim - 2 is the spatial rank. + # Like every validator in this module this relies on meta["val"] and fails open + # when it is absent. + if not args_bounds_check(node.args, _CONV_ARG_TRANSPOSED): input_node = node.args[_CONV_ARG_INPUT] val = input_node.meta.get("val") if hasattr(input_node, "meta") else None - if (ndim := getattr(val, "ndim", None)) == 5: - _LOGGER.debug( - "3D convolution '%s' (ndim %s) is not supported on TensorRT-RTX for " - "Turing (SM 7.5). Falling back to PyTorch.", - node.name, - ndim, - ) + spatial_rank = val.ndim - 2 if val is not None else None + if turing_rejects_forward_convolution(node, spatial_rank, settings): return False return True diff --git a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py index 3f15c25cad..9cc8532929 100644 --- a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py @@ -1,17 +1,22 @@ # mypy: disallow-untyped-decorators=False import logging -from typing import Dict, Sequence, Tuple, Union +from typing import Dict, Optional, Sequence, Tuple, Union import tensorrt as trt -from torch.fx.node import Argument, Target +from torch.fx.node import Argument, Node, Target from torch_tensorrt._features import ENABLED_FEATURES +from torch_tensorrt.dynamo._settings import CompilationSettings 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._ConverterRegistry import ( dynamo_tensorrt_converter, ) +from torch_tensorrt.dynamo.conversion.aten_ops_converters import ( + turing_rejects_forward_convolution, +) +from torch_tensorrt.dynamo.conversion.converter_utils import args_bounds_check from torch_tensorrt.dynamo.lowering.passes.fuse_distributed_ops import ( tensorrt_fused_nccl_all_gather_op, tensorrt_fused_nccl_all_reduce_op, @@ -27,7 +32,34 @@ _LOGGER: logging.Logger = logging.getLogger(__name__) -@dynamo_tensorrt_converter(tensorrt_conv_asym_pad_op, supports_dynamic_shapes=True) +def conv_asym_pad_capability_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Apply the aten.convolution capability guards to the fused pad+conv op. + + ``fuse_pad_into_convolution`` rewrites constant_pad_nd -> aten.convolution into + this op during lowering, which erases the aten.convolution node that + ``convolution_capability_validator`` would otherwise have rejected. Without this, + the pass smuggles a 3D convolution past that guard onto Turing (SM 7.5), where it + surfaces at runtime as a null execution context. The pass runs in the normal + compile path, so this is not confined to tests. + + Spatial rank comes from the argument list rather than node meta, so the guard holds + whether or not the caller traced with dynamo. + """ + # conv_asym_pad(source, weight, bias, stride, pre_padding, post_padding, dilation, + # groups) -- no transposed argument; the pass skips those outright. + _ARG_STRIDE = 3 + stride = args_bounds_check(node.args, _ARG_STRIDE) + spatial_rank = len(stride) if stride is not None else None + return not turing_rejects_forward_convolution(node, spatial_rank, settings) + + +@dynamo_tensorrt_converter( + tensorrt_conv_asym_pad_op, + capability_validator=conv_asym_pad_capability_validator, + supports_dynamic_shapes=True, +) def conv_asym_pad( ctx: ConversionContext, target: Target, diff --git a/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py index f023aad984..a909cb4ffa 100644 --- a/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py +++ b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py @@ -5,7 +5,7 @@ import torch.nn.functional as F from parameterized import parameterized -from ..conversion.harness import DispatchTestCase +from ..conversion.harness import DispatchTestCase, skip_if_trt_rtx_turing def _node_targets(gm: torch.fx.GraphModule) -> list: @@ -230,6 +230,11 @@ class TestFusePadIntoConvolutionConverter(DispatchTestCase): ] ) def test_padded_conv3d(self, _, pad, stride, dilation, groups, bias): + # The fusion turns this into tensorrt::conv_asym_pad, whose validator rejects + # 3D on Turing. DispatchTestCase has no PyTorch-fallback path, so without this + # the rejection surfaces as UnsupportedOperatorException. + skip_if_trt_rtx_turing(self, "3D convolution") + class PaddedConv3d(nn.Module): def __init__(self) -> None: super().__init__() diff --git a/tests/py/dynamo/models/test_turing_capability_guards.py b/tests/py/dynamo/models/test_turing_capability_guards.py index e7f5044f40..023bbe32d8 100644 --- a/tests/py/dynamo/models/test_turing_capability_guards.py +++ b/tests/py/dynamo/models/test_turing_capability_guards.py @@ -70,6 +70,33 @@ def forward(self, x): return self.conv(x) +class PaddedConv3d(nn.Module): + """A conv3d the 3D-conv guard never sees as a convolution. + + fuse_pad_into_convolution folds the pad into the conv and rewrites the pair as + tensorrt::conv_asym_pad, erasing the aten.convolution node. It fires for any + zero-fill, non-negative pad -- asymmetry is not required. + """ + + def __init__(self): + super().__init__() + self.conv = nn.Conv3d(4, 8, 3, padding=0) + + def forward(self, x): + return self.conv(F.pad(x, (1, 1, 1, 1, 2, 0))) + + +class PaddedConv2d(nn.Module): + """Same fusion, 2D -- which Turing does support.""" + + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(4, 8, 3, padding=0) + + def forward(self, x): + return self.conv(F.pad(x, (0, 1, 0, 1))) + + @unittest.skipIf( not ENABLED_FEATURES.tensorrt_rtx, "Turing capability guards only apply to TensorRT-RTX", @@ -154,6 +181,23 @@ def test_declared_turing_target_keeps_conv2d_on_trt(self): compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) self.assertGreater(_trt_submodule_count(compiled), 0) + def test_declared_turing_target_falls_back_padded_conv3d(self): + # The fused op needs its own guard: by partitioning time the aten.convolution + # node the conv guard keys on no longer exists. + mod = PaddedConv3d() + inputs = (torch.randn(1, 4, 8, 8, 8).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertEqual(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + + def test_declared_turing_target_keeps_padded_conv2d_on_trt(self): + # Only rank 3 is rejected; the fused op must keep serving 2D. + mod = PaddedConv2d() + inputs = (torch.randn(1, 4, 16, 16).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertGreater(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + def test_declared_turing_target_keeps_transposed_conv3d_on_trt(self): # Transposed 3D convolution is a distinct layer and does work on Turing. mod = ConvTranspose3d()