From 0172589da933b01b2dfb3ea9f34dc473ac711108 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Sat, 29 Aug 2026 17:38:57 -0700 Subject: [PATCH 1/3] feat: target TensorRT-RTX compute capabilities explicitly Adds a target_compute_capabilities setting naming the architectures an engine is built for, and declares them on the TensorRT-RTX builder config, so an artifact can be produced for something other than the build host. Declaring it on standard TensorRT raises: that backend always builds for the current device. An undeclared target is now resolved rather than skipped. Left at the builder default num_compute_capabilities stays 0, and a refittable graph then fails with "Compatible cubin or ptx module for device target '75' not found" -- so on Turing every refittable build failed. It now resolves to the current device the way partitioning already did via get_target_compute_capabilities(), naming ComputeCapability.CURRENT rather than an SM lookup because TensorRT-RTX names only a subset of architectures. The builder-config work lives in set_rtx_compute_capabilities() so it can be unit tested against a stub config, with no GPU and no engine build. Not exposed on cross_compile_for_windows: TensorRT-RTX does not support that path. The setting is engine-invariant: an engine built for one capability set must not be reused for a compile targeting another. --- py/torch_tensorrt/_utils.py | 34 ++++++++- py/torch_tensorrt/dynamo/_compiler.py | 10 +++ py/torch_tensorrt/dynamo/_defaults.py | 1 + py/torch_tensorrt/dynamo/_settings.py | 20 ++++- .../dynamo/conversion/_TRTInterpreter.py | 75 ++++++++++++++++++- 5 files changed, 135 insertions(+), 5 deletions(-) diff --git a/py/torch_tensorrt/_utils.py b/py/torch_tensorrt/_utils.py index 5d43db0e57..d1b82d7d42 100644 --- a/py/torch_tensorrt/_utils.py +++ b/py/torch_tensorrt/_utils.py @@ -7,7 +7,7 @@ import tempfile import urllib.request from pathlib import Path -from typing import Any, Optional +from typing import Any, Optional, Tuple import tensorrt as trt import torch @@ -380,3 +380,35 @@ def load_tensorrt_llm_for_nccl() -> bool: plugin_lib_path = download_and_get_plugin_lib_path() return load_and_initialize_trtllm_plugin(plugin_lib_path) # type: ignore[arg-type] return False + + +# --- TensorRT-RTX architecture targeting ------------------------------------- + +TURING_COMPUTE_CAPABILITY = (7, 5) + + +def get_target_compute_capabilities( + settings: Optional[Any] = None, +) -> Tuple[Tuple[int, int], ...]: + """Compute capabilities this compilation targets: the declared targets, or the + current device when none were declared. + + Capability validators key off this rather than the build host, so a module compiled + on Ampere and shipped to Turing does not retain ops Turing cannot execute. + """ + if settings and (targets := getattr(settings, "target_compute_capabilities", None)): + return tuple((int(major), int(minor)) for major, minor in targets) + return (torch.cuda.get_device_capability(),) + + +def trt_rtx_targets_turing(settings: Optional[Any] = None) -> bool: + """True when TensorRT-RTX is in use and SM 7.5 is among the build targets. + + A single compiled artifact carries a single partitioning, so an op unsupported on + *any* targeted architecture must fall back to PyTorch for all of them. + """ + from torch_tensorrt._features import ENABLED_FEATURES + + if not ENABLED_FEATURES.tensorrt_rtx: + return False + return TURING_COMPUTE_CAPABILITY in get_target_compute_capabilities(settings) diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index f2f6423e41..ad9d9017f5 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -466,6 +466,9 @@ def compile( enable_experimental_decompositions: bool = _defaults.ENABLE_EXPERIMENTAL_DECOMPOSITIONS, dryrun: bool = _defaults.DRYRUN, hardware_compatible: bool = _defaults.HARDWARE_COMPATIBLE, + target_compute_capabilities: Optional[ + List[Tuple[int, int]] + ] = _defaults.TARGET_COMPUTE_CAPABILITIES, timing_cache_path: str = _defaults.TIMING_CACHE_PATH, lazy_engine_init: bool = _defaults.LAZY_ENGINE_INIT, cache_built_engines: bool = _defaults.CACHE_BUILT_ENGINES, @@ -564,6 +567,7 @@ def compile( enable_experimental_decompositions (bool): Use the full set of operator decompositions. These decompositions may not be tested but serve to make the graph easier to convert to TensorRT, potentially increasing the amount of graphs run in TensorRT. dryrun (bool): Toggle for "Dryrun" mode, running everything except conversion to TRT and logging outputs hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer) + target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only. Drives both engine targeting and op partitioning, so ops unsupported on any listed target fall back to PyTorch. timing_cache_path (str): Path to the timing cache if it exists (or) where it will be saved after compilation. Not used for TensorRT-RTX. lazy_engine_init (bool): Defer setting up engines until the compilation of all engines is complete. Can allow larger models with multiple graph breaks to compile but can lead to oversubscription of GPU memory at runtime. cache_built_engines (bool): Whether to save the compiled TRT engines to storage @@ -767,6 +771,7 @@ def compile( "dla_global_dram_size": dla_global_dram_size, "dryrun": dryrun, "hardware_compatible": hardware_compatible, + "target_compute_capabilities": target_compute_capabilities, "timing_cache_path": timing_cache_path, "lazy_engine_init": lazy_engine_init, "cache_built_engines": cache_built_engines, @@ -1795,6 +1800,9 @@ def convert_exported_program_to_serialized_trt_engine( enable_experimental_decompositions: bool = _defaults.ENABLE_EXPERIMENTAL_DECOMPOSITIONS, dryrun: bool = _defaults.DRYRUN, hardware_compatible: bool = _defaults.HARDWARE_COMPATIBLE, + target_compute_capabilities: Optional[ + List[Tuple[int, int]] + ] = _defaults.TARGET_COMPUTE_CAPABILITIES, timing_cache_path: str = _defaults.TIMING_CACHE_PATH, lazy_engine_init: bool = _defaults.LAZY_ENGINE_INIT, cache_built_engines: bool = _defaults.CACHE_BUILT_ENGINES, @@ -1897,6 +1905,7 @@ def convert_exported_program_to_serialized_trt_engine( enable_experimental_decompositions (bool): Use the full set of operator decompositions. These decompositions may not be tested but serve to make the graph easier to convert to TensorRT, potentially increasing the amount of graphs run in TensorRT. dryrun (bool): Toggle for "Dryrun" mode, running everything except conversion to TRT and logging outputs hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer) + target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only. Drives both engine targeting and op partitioning, so ops unsupported on any listed target fall back to PyTorch. timing_cache_path (str): Path to the timing cache if it exists (or) where it will be saved after compilation. Not used for TensorRT-RTX. lazy_engine_init (bool): Defer setting up engines until the compilation of all engines is complete. Can allow larger models with multiple graph breaks to compile but can lead to oversubscription of GPU memory at runtime. cache_built_engines (bool): Whether to save the compiled TRT engines to storage @@ -2081,6 +2090,7 @@ def convert_exported_program_to_serialized_trt_engine( "dla_global_dram_size": dla_global_dram_size, "dryrun": dryrun, "hardware_compatible": hardware_compatible, + "target_compute_capabilities": target_compute_capabilities, "timing_cache_path": timing_cache_path, "lazy_engine_init": lazy_engine_init, "cache_built_engines": cache_built_engines, diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py index 1d6c65dd6f..66a7d7eea5 100644 --- a/py/torch_tensorrt/dynamo/_defaults.py +++ b/py/torch_tensorrt/dynamo/_defaults.py @@ -63,6 +63,7 @@ ENABLE_CROSS_COMPILE_FOR_WINDOWS = False TILING_OPTIMIZATION_LEVEL = "none" L2_LIMIT_FOR_TILING = -1 +TARGET_COMPUTE_CAPABILITIES = None USE_DISTRIBUTED_MODE_TRACE = False OFFLOAD_MODULE_TO_CPU = False ENABLE_AUTOCAST = False diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index a030f081b6..e7f57c8967 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import Any, Collection, Optional, Set, Tuple, Union +from typing import Any, Collection, List, Optional, Set, Tuple, Union import tensorrt as trt import torch @@ -47,6 +47,7 @@ REUSE_CACHED_ENGINES, SPARSE_WEIGHTS, STRIP_ENGINE_WEIGHTS, + TARGET_COMPUTE_CAPABILITIES, TILING_OPTIMIZATION_LEVEL, TIMING_CACHE_PATH, TRUNCATE_DOUBLE, @@ -71,7 +72,9 @@ def _normalize_disabled_constant_fold_exclusions( validate_disabled_constant_fold_exclusions, ) - return validate_disabled_constant_fold_exclusions(rule_ids) + # Typed local: --strict mypy rejects returning the untyped helper directly. + normalized: Set[str] = validate_disabled_constant_fold_exclusions(rule_ids) + return normalized @dataclass @@ -106,6 +109,7 @@ class CompilationSettings: TRT Engines. Prints detailed logs of the graph structure and nature of partitioning. Optionally saves the output to a file if a string path is specified hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer) + target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only -- raises on standard TensorRT. Drives both engine targeting and op partitioning: a compiled artifact carries a single partitioning, so an op unsupported on any listed target falls back to PyTorch for all of them. timing_cache_path (str): Path to the timing cache if it exists (or) where it will be saved after compilation. Not used for TensorRT-RTX (no autotuning). cache_built_engines (bool): Whether to save the compiled TRT engines to storage reuse_cached_engines (bool): Whether to load the compiled TRT engines from storage @@ -181,6 +185,9 @@ class CompilationSettings: enable_cross_compile_for_windows: bool = ENABLE_CROSS_COMPILE_FOR_WINDOWS tiling_optimization_level: str = TILING_OPTIMIZATION_LEVEL l2_limit_for_tiling: int = L2_LIMIT_FOR_TILING + target_compute_capabilities: Optional[List[Tuple[int, int]]] = ( + TARGET_COMPUTE_CAPABILITIES + ) use_distributed_mode_trace: bool = USE_DISTRIBUTED_MODE_TRACE offload_module_to_cpu: bool = OFFLOAD_MODULE_TO_CPU enable_autocast: bool = ENABLE_AUTOCAST @@ -212,6 +219,14 @@ def __post_init__(self) -> None: self.disabled_constant_fold_exclusions ) ) + if self.target_compute_capabilities: + from torch_tensorrt._features import ENABLED_FEATURES + + if not ENABLED_FEATURES.tensorrt_rtx: + raise ValueError( + "target_compute_capabilities is only supported on TensorRT-RTX. " + "Standard TensorRT always builds for the current device." + ) def __getstate__(self) -> dict[str, Any]: from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( @@ -245,6 +260,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: "sparse_weights", "engine_capability", "hardware_compatible", + "target_compute_capabilities", "refit_identical_engine_weights", "immutable_weights", "enable_weight_streaming", diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py index 448d96a0ce..fe4f79d9af 100644 --- a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py +++ b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py @@ -15,6 +15,7 @@ ) import numpy as np +import tensorrt as trt import torch import torch.fx from torch.fx.experimental.proxy_tensor import unset_fake_temporarily @@ -52,8 +53,6 @@ ) from torch_tensorrt.logging import TRT_LOGGER -import tensorrt as trt - _LOGGER: logging.Logger = logging.getLogger(__name__) TRT_INTERPRETER_CALL_PRE_OBSERVER: Observer[Callable[[torch.fx.GraphModule], None]] = ( @@ -86,6 +85,73 @@ class TRTInterpreterResult(NamedTuple): aliased_io: dict[str, tuple[str, str]] = {} +def _named_compute_capability(major: int, minor: int) -> "trt.ComputeCapability": + """Map a compute capability to the TensorRT-RTX enum member naming it. + + The return annotation is quoted: ``ComputeCapability`` exists only in the + TensorRT-RTX bindings, and an unquoted annotation is evaluated at import time on + every backend. + + Raises when TensorRT-RTX does not name the architecture: an explicit request for an + unsupported target is a user error and should stay loud. + """ + name = f"SM{major}{minor}" + if (compute_capability := getattr(trt.ComputeCapability, name, None)) is None: + supported = [m for m in dir(trt.ComputeCapability) if m.startswith("SM")] + raise ValueError( + f"TensorRT-RTX has no compute capability {name} for requested " + f"target ({major}, {minor}). Supported: {supported}" + ) + return compute_capability + + +def _set_declared_compute_capabilities( + builder_config: trt.IBuilderConfig, + declared: Sequence[Tuple[int, int]], +) -> None: + """Declare the architectures the artifact is built for, by name.""" + builder_config.num_compute_capabilities = len(declared) + for idx, (major, minor) in enumerate(declared): + if not builder_config.set_compute_capability( + _named_compute_capability(major, minor), idx + ): + raise RuntimeError( + f"Failed to set TensorRT-RTX compute capability SM{major}{minor}" + ) + _LOGGER.info(f"Targeting TensorRT-RTX compute capabilities {declared}") + + +def _set_current_compute_capability(builder_config: trt.IBuilderConfig) -> None: + """Declare the current device, for a build that named no targets. + + ComputeCapability.CURRENT rather than an SM lookup: TensorRT-RTX names + only a subset of architectures, and the implicit path must not start failing on + hosts that build fine today. + """ + builder_config.num_compute_capabilities = 1 + if not builder_config.set_compute_capability(trt.ComputeCapability.CURRENT, 0): + raise RuntimeError( + "Failed to set the TensorRT-RTX compute capability of the current device" + ) + _LOGGER.info("Targeting the TensorRT-RTX compute capability of the current device") + + +def set_rtx_compute_capabilities( + builder_config: trt.IBuilderConfig, + declared: Optional[Sequence[Tuple[int, int]]], +) -> None: + """Tell the TensorRT-RTX builder which architectures to build for. + + An undeclared target is resolved rather than skipped: the builder default leaves + num_compute_capabilities at 0, and a refittable graph then fails with "Compatible + cubin or ptx module for device target '75' not found". + """ + if declared: + _set_declared_compute_capabilities(builder_config, declared) + else: + _set_current_compute_capability(builder_config) + + @cls_supports_debugger class TRTInterpreter(torch.fx.Interpreter): # type: ignore[misc] def __init__( @@ -385,6 +451,11 @@ def _populate_trt_builder_config( self.compilation_settings.l2_limit_for_tiling ) + if ENABLED_FEATURES.tensorrt_rtx: + set_rtx_compute_capabilities( + builder_config, self.compilation_settings.target_compute_capabilities + ) + return builder_config def _create_timing_cache( From 100ef2a00bd066db742fb33c4e95b87f8a64da60 Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Sat, 29 Aug 2026 15:58:35 -0700 Subject: [PATCH 2/3] fix: fall back to PyTorch for ops TensorRT-RTX cannot serve on Turing TensorRT-RTX does not support FP32 GEMMs or 3D convolutions on Turing (SM 7.5), and Turing has no bfloat16 hardware. Handed those ops anyway it returns a null execution context, segfaults on bf16, or -- under dynamic shapes -- builds, runs and returns an all-zero tensor with no exception. That silent case is the motivating one. Guards key off the capabilities being built for rather than the build host, so an ahead-of-time build for another architecture partitions correctly. The GEMM guard keys on fp32 operands only; the convolution guard covers forward 3D only, since transposed 3D works on Turing. bfloat16 is gated in the partitioners because the crash is not operator-specific, mirroring the existing complex-dtype handling. Converter unit tests need explicit skips: DispatchTestCase bypasses the partitioner, so a guarded node raises UnsupportedOperatorException instead of falling back and an unguarded one reaches TensorRT-RTX and fails. The cdist skips state the converter's condition rather than p == 2, because a GEMM is only emitted for compute_mode 1, or 0/absent with an operand above the row threshold. Measured on a T4: the 3 cases outside that condition pass, the 9 inside it fail. Rather than hand-copy the branch, it is extracted from cdist_forward as cdist_emits_matmul/CDIST_MATMUL_ROW_THRESHOLD -- the converter now consumes its own predicate, and the test and the Turing capability validator consume the same one, so none of the three can drift. The threshold was a bare literal in the converter until now. Four parameterisations named for a compute_mode they do not use are renamed. Also clears two pre-existing lint failures in aten_ops_converters.py that pre-commit blocks on now the file is in the changed set: an unused type: ignore[assignment] and a "mis-evaluate" spelling. --- .../dynamo/conversion/aten_ops_converters.py | 103 ++++++++++++++++-- .../conversion/impl/normalization/ops.py | 32 +++++- .../partitioning/_adjacency_partitioner.py | 9 ++ .../partitioning/_global_partitioner.py | 31 ++++++ tests/py/dynamo/conversion/harness.py | 21 ++++ tests/py/dynamo/conversion/test_addmm_aten.py | 6 +- .../dynamo/conversion/test_binary_ops_aten.py | 4 +- tests/py/dynamo/conversion/test_casts.py | 4 +- tests/py/dynamo/conversion/test_cat_aten.py | 9 +- tests/py/dynamo/conversion/test_cdist_aten.py | 25 ++++- .../conversion/test_convolution_aten.py | 6 +- .../dynamo/conversion/test_index_put_aten.py | 8 +- .../py/dynamo/conversion/test_matmul_aten.py | 12 +- tests/py/dynamo/conversion/test_where_aten.py | 4 +- 14 files changed, 245 insertions(+), 29 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 4f245a3883..cc679ec6c2 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -24,6 +24,7 @@ from torch_tensorrt._utils import ( is_tensorrt_rtx_version_supported, is_tensorrt_version_supported, + trt_rtx_targets_turing, ) from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo._SourceIR import SourceIR @@ -846,12 +847,62 @@ def aten_ops_gelu( ) -@dynamo_tensorrt_converter(torch.ops.aten.matmul, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.matmul.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.dot.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.mm.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.mv.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.bmm.default, supports_dynamic_shapes=True) +def gemm_capability_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Reject FP32 GEMMs on TensorRT-RTX when Turing (SM 7.5) is a build target. + + Only operand dtype matters, so fp16 GEMMs accumulating in fp32 (``use_fp32_acc``) + keep running on TensorRT. + """ + if not trt_rtx_targets_turing(settings): + return True + + def is_fp32(operand: Argument) -> bool: + val = operand.meta.get("val") if hasattr(operand, "meta") else None + return bool(getattr(val, "dtype", None) == torch.float32) + + if any(map(is_fp32, node.args[:2])): + _LOGGER.debug( + "FP32 GEMM '%s' is not supported on TensorRT-RTX for Turing " + "(SM 7.5). Falling back to PyTorch.", + node.name, + ) + return False + + return True + + +@dynamo_tensorrt_converter( + torch.ops.aten.matmul, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.matmul.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.dot.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.mm.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.mv.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.bmm.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) def aten_ops_matmul( ctx: ConversionContext, target: Target, @@ -1227,7 +1278,7 @@ def _index_copy_kv_eligible( if len(node.args) < 4: return False if input_node is None: - input_node = node.args[0] # type: ignore[assignment] + input_node = node.args[0] dim, _index_node, src_node = node.args[1:4] if not isinstance(input_node, Node) or input_node.op != "placeholder": @@ -2838,7 +2889,7 @@ def aten_ops_logical_xor( def _is_absorbing_bitwise_scalar(target: Target, scalar: Any) -> bool: - """Would TensorRT mis-evaluate this scalar bitwise op (see above)?""" + """Would TensorRT evaluate this scalar bitwise op incorrectly (see above)?""" if target not in _ABSORBING_BITWISE_SCALAR: return False # A non-bool scalar is rejected by the dtype check further down anyway, and @@ -3193,6 +3244,14 @@ def aten_ops_le( ) +# aten.convolution(input, weight, bias, stride, padding, dilation, transposed, +# output_padding, groups) +_CONV_ARG_INPUT = 0 +_CONV_ARG_STRIDE = 3 +_CONV_ARG_DILATION = 5 +_CONV_ARG_TRANSPOSED = 6 + + def convolution_capability_validator( node: Node, settings: Optional[CompilationSettings] = None ) -> bool: @@ -3206,9 +3265,9 @@ def convolution_capability_validator( return True if ( - args_bounds_check(node.args, 6) # transposed? - and (stride := args_bounds_check(node.args, 3)) - and (dilation := args_bounds_check(node.args, 5)) + args_bounds_check(node.args, _CONV_ARG_TRANSPOSED) + and (stride := args_bounds_check(node.args, _CONV_ARG_STRIDE)) + and (dilation := args_bounds_check(node.args, _CONV_ARG_DILATION)) and any(s > 1 for s in stride) and any(d > 1 for d in dilation) ): @@ -3219,6 +3278,22 @@ 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: + 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, + ) + return False + return True @@ -3654,7 +3729,11 @@ def aten_ops_argmin( ) -@dynamo_tensorrt_converter(torch.ops.aten.addmm.default, supports_dynamic_shapes=True) +@dynamo_tensorrt_converter( + torch.ops.aten.addmm.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) @enforce_tensor_types( { 0: (TRTTensor,), diff --git a/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py index a42cf8ff43..435c688778 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py @@ -643,6 +643,31 @@ def tri_upper_indices( return indices_tensor.get_output(0) +#: Above this many rows in either operand, cdist_forward computes p == 2 with a +#: matrix-multiply layer instead of a broadcast subtract. +CDIST_MATMUL_ROW_THRESHOLD = 25 + + +def cdist_emits_matmul( + p: float, + compute_mode: Optional[int], + rows: Sequence[Optional[int]], +) -> bool: + """Would :func:`cdist_forward` emit a matrix-multiply layer for these arguments? + + ``rows`` is ``shape[-2]`` of each operand, or ``None`` where it is not statically + known; an unknown row count is treated as below the threshold, so callers that + cannot see shapes fail open. + """ + if p != 2: + return False + if compute_mode == 1: + return True + if compute_mode in (0, None): + return any(r is not None and r > CDIST_MATMUL_ROW_THRESHOLD for r in rows) + return False + + def cdist_forward( ctx: ConversionContext, target: Target, @@ -669,7 +694,8 @@ def cdist_forward( p (float): p value for the p-norm distance to calculate between each vector pair compute_mode (int): Controls the computation method based on the size of the input sets: - None ('use_mm_for_euclid_dist_if_necessary'): Default mode. Uses matrix multiplication to calculate - Euclidean distance (p=2) if either the number of vectors in x1 or x2 exceeds 25 (P > 25 or R > 25). + Euclidean distance (p=2) if the number of vectors in x1 or x2 exceeds + CDIST_MATMUL_ROW_THRESHOLD. - 1 ('use_mm_for_euclid_dist'): Always use matrix multiplication approach to calculate euclidean distance (p = 2) - 2 ('donot_use_mm_for_euclid_dist'): Never use matrix multiplication approach to calculate @@ -725,9 +751,7 @@ def cdist_forward( ctx, target, source_ir, f"{name}_sum", abs_val, dim=-1, keepdim=False ) elif p == 2: - if ( - compute_mode == 0 and (x1.shape[-2] > 25 or x2.shape[-2] > 25) - ) or compute_mode == 1: + if cdist_emits_matmul(p, compute_mode, (x1.shape[-2], x2.shape[-2])): # Compute squared elements x1_squared = impl.elementwise.pow( ctx, target, source_ir, f"{name}_x1_squared", x1, 2 diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 098e9b2685..2565b18bdd 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -54,6 +54,15 @@ def is_node_supported( return False settings = CONVERTERS.compilation_settings + if TorchTensorRTOperatorSupport._has_bf16_on_turing(node, settings): + # bfloat16 has no Turing hardware; compiling it for SM 7.5 crashes the + # process, so force the PyTorch fallback. + if not node.is_impure(): + self.unsupported_operators[node_name] = ( + self.unsupported_operators.get(node_name, 0) + 1 + ) + return False + if ( settings is not None and settings.fallback_data_dependent_ops diff --git a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py index 68e35e060a..2b1a3db4d0 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py @@ -6,6 +6,7 @@ from torch.fx.node import Target from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition from torch.fx.passes.operator_support import OperatorSupport, SupportDict +from torch_tensorrt._utils import trt_rtx_targets_turing from torch_tensorrt.dynamo._defaults import ( MIN_BLOCK_SIZE, REQUIRE_FULL_COMPILATION, @@ -166,6 +167,27 @@ def _dtype(n: torch.fx.Node) -> Optional[torch.dtype]: return True return False + @staticmethod + def _has_bf16_on_turing( + node: torch.fx.Node, settings: Optional[object] = None + ) -> bool: + """Return True if this node touches bfloat16 while targeting Turing (SM 7.5). + + Turing has no bfloat16 hardware, and compiling a bf16 node for SM 7.5 is an + error, so these nodes must run in the PyTorch fallback. Checked at the + partitioner rather than per-converter because it is not operator-specific. + """ + if not trt_rtx_targets_turing(settings): + return False + + def _dtype(n: torch.fx.Node) -> Optional[torch.dtype]: + val = n.meta.get("val") + return getattr(val, "dtype", None) if val is not None else None + + if _dtype(node) == torch.bfloat16: + return True + return any(_dtype(arg) == torch.bfloat16 for arg in node.all_input_nodes) + @staticmethod def _requires_output_allocator(node: torch.fx.Node) -> bool: # True if the converter selected for this node needs a TRT output allocator, @@ -191,6 +213,15 @@ def is_node_supported( return False settings = CONVERTERS.compilation_settings + if self._has_bf16_on_turing(node, settings): + # bfloat16 has no Turing hardware; compiling it for SM 7.5 crashes the + # process, so force the PyTorch fallback. + if not node.is_impure(): + self.unsupported_operators[node_name] = ( + self.unsupported_operators.get(node_name, 0) + 1 + ) + return False + if ( settings is not None and settings.fallback_data_dependent_ops diff --git a/tests/py/dynamo/conversion/harness.py b/tests/py/dynamo/conversion/harness.py index 2dab3d64df..81d834f7e3 100644 --- a/tests/py/dynamo/conversion/harness.py +++ b/tests/py/dynamo/conversion/harness.py @@ -89,6 +89,27 @@ def infer_module_output_dtypes_for_test( # this is to enable dynamo tracer as True in the converter test files batch by batch +def skip_if_trt_rtx_turing(test_case: TestCase, what: str) -> None: + """Skip a converter test for a case TensorRT-RTX cannot serve on Turing (SM 7.5). + + ``DispatchTestCase.run_test`` hands the graph straight to ``TRTInterpreter``, + skipping the partitioner, so there is no PyTorch fallback here: a guarded node + raises ``UnsupportedOperatorException`` and an unguarded one reaches TensorRT-RTX + and fails, or for bf16 crashes the process. Either way the test needs an explicit + skip. + """ + import torch_tensorrt + + if ( + torch_tensorrt.ENABLED_FEATURES.tensorrt_rtx + and torch.cuda.is_available() + and torch.cuda.get_device_capability() == (7, 5) + ): + test_case.skipTest( + f"{what} is not supported on TensorRT-RTX for Turing (SM 7.5)" + ) + + def get_use_dynamo_tracer(use_dynamo_tracer: Any) -> bool: # if in our converter tests we specifically set use_dynamo_tracer field, honor it if use_dynamo_tracer is not None and isinstance(use_dynamo_tracer, bool): diff --git a/tests/py/dynamo/conversion/test_addmm_aten.py b/tests/py/dynamo/conversion/test_addmm_aten.py index 6108d3ea6d..7ffaf279b1 100644 --- a/tests/py/dynamo/conversion/test_addmm_aten.py +++ b/tests/py/dynamo/conversion/test_addmm_aten.py @@ -3,7 +3,7 @@ from parameterized import parameterized from torch.testing._internal.common_utils import run_tests -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestAddmmConverter(DispatchTestCase): @@ -17,6 +17,8 @@ class TestAddmmConverter(DispatchTestCase): ] ) def test_addmm(self, input_shape, mat1_shape, mat2_shape): + skip_if_trt_rtx_turing(self, "aten.addmm (an FP32 GEMM)") + class Addmm(nn.Module): def forward(self, input, mat1, mat2): return torch.ops.aten.addmm.default(input, mat1, mat2) @@ -43,6 +45,8 @@ def forward(self, input, mat1, mat2): ] ) def test_addmm_scale(self, input_shape, mat1_shape, mat2_shape, beta, alpha): + skip_if_trt_rtx_turing(self, "aten.addmm (an FP32 GEMM)") + class Addmm(nn.Module): def forward(self, input, mat1, mat2): return torch.ops.aten.addmm.default( diff --git a/tests/py/dynamo/conversion/test_binary_ops_aten.py b/tests/py/dynamo/conversion/test_binary_ops_aten.py index 4ac613adea..8a288ab5ff 100644 --- a/tests/py/dynamo/conversion/test_binary_ops_aten.py +++ b/tests/py/dynamo/conversion/test_binary_ops_aten.py @@ -8,7 +8,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing NEED_TEST_BOTH_CONSTANTS_CASE = True @@ -238,6 +238,8 @@ def forward(self, x, y): ] ) def test_elementwise_ops_bf16(self, _, orig_op): + skip_if_trt_rtx_turing(self, "bfloat16") + class TestModule(nn.Module): def __init__(self, orig_op): super().__init__() diff --git a/tests/py/dynamo/conversion/test_casts.py b/tests/py/dynamo/conversion/test_casts.py index 550a9c1d45..8f5a335aad 100644 --- a/tests/py/dynamo/conversion/test_casts.py +++ b/tests/py/dynamo/conversion/test_casts.py @@ -9,7 +9,7 @@ from torch_tensorrt import dtype from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestCloneConverter(DispatchTestCase): @@ -68,6 +68,8 @@ def forward(self, x): ) def test_to_copy_bfloat16(self): + skip_if_trt_rtx_turing(self, "bfloat16") + class ToCopyBFloat16(nn.Module): def forward(self, x): y = torch.ops.aten._to_copy.default(x, dtype=torch.bfloat16) diff --git a/tests/py/dynamo/conversion/test_cat_aten.py b/tests/py/dynamo/conversion/test_cat_aten.py index 475127210f..150e72e04f 100644 --- a/tests/py/dynamo/conversion/test_cat_aten.py +++ b/tests/py/dynamo/conversion/test_cat_aten.py @@ -4,7 +4,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestCatConverter(DispatchTestCase): @@ -177,6 +177,10 @@ def forward(self, x): def test_cat_three_different_dtypes(self): """Test cat with three different dtypes - bfloat16, float16, float32""" + # The cat promotes bf16 away, but TensorRT-RTX still rejects the network for + # naming the type: "Cannot compile for target(s) sm75 ... HW-specific + # datatypes: b16". That only surfaces once the capability is declared. + skip_if_trt_rtx_turing(self, "bfloat16") class ThreeDtypeCat(nn.Module): def __init__(self): @@ -344,6 +348,9 @@ def forward(self, x): def test_cat_bf16_dtype_preservation(self): """Test that bfloat16 dtype is preserved in constant layers (not converted to fp32)""" + # The engine's input and output are both bf16 here, so bf16 really does reach + # TensorRT. + skip_if_trt_rtx_turing(self, "bfloat16") class CatBF16Constants(nn.Module): def __init__(self): diff --git a/tests/py/dynamo/conversion/test_cdist_aten.py b/tests/py/dynamo/conversion/test_cdist_aten.py index 71628df510..719d739d9b 100644 --- a/tests/py/dynamo/conversion/test_cdist_aten.py +++ b/tests/py/dynamo/conversion/test_cdist_aten.py @@ -2,8 +2,11 @@ import torch.nn as nn from parameterized import parameterized from torch.testing._internal.common_utils import run_tests +from torch_tensorrt.dynamo.conversion.impl.normalization.ops import ( + cdist_emits_matmul, +) -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestCdistConverter(DispatchTestCase): @@ -21,6 +24,9 @@ class TestCdistConverter(DispatchTestCase): ] ) def test_cdist_float_same_shape(self, name, shape, p, compute_mode): + if cdist_emits_matmul(p, compute_mode, (shape[-2], shape[-2])): + skip_if_trt_rtx_turing(self, "cdist whose converter emits a GEMM") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) @@ -47,6 +53,9 @@ def forward(self, x1, x2): def test_cdist_float_broadcast_and_diff_shape( self, name, shape_1, shape_2, p, compute_mode ): + if cdist_emits_matmul(p, compute_mode, (shape_1[-2], shape_2[-2])): + skip_if_trt_rtx_turing(self, "cdist whose converter emits a GEMM") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) @@ -59,16 +68,19 @@ def forward(self, x1, x2): @parameterized.expand( [ - ("compute_mode_0", (15, 10, 5), (15, 35, 5), 2.0, 0), - ("compute_mode_1", (35, 35, 5), (35, 45, 5), 2.0, 0), - ("compute_mode_2", (15, 10, 5), (15, 35, 5), 2.0, 1), - ("compute_mode_3", (35, 35, 5), (35, 45, 5), 2.0, 2), + ("mode_0_one_operand_above_threshold", (15, 10, 5), (15, 35, 5), 2.0, 0), + ("mode_0_both_operands_above_threshold", (35, 35, 5), (35, 45, 5), 2.0, 0), + ("mode_1_always_matmul", (15, 10, 5), (15, 35, 5), 2.0, 1), + ("mode_2_never_matmul", (35, 35, 5), (35, 45, 5), 2.0, 2), ("p_2_mm_shape_1", (2, 2, 14, 5), (3, 5), 2, 1), ("p_2_mm_shape_2", (2, 2, 14, 5), (2, 3, 5), 2, 1), ("p_2_mm_shape_3", (2, 2, 14, 5), (2, 2, 3, 5), 2, 1), ] ) def test_cdist_p_2_compute_mode(self, name, shape_1, shape_2, p, compute_mode): + if cdist_emits_matmul(p, compute_mode, (shape_1[-2], shape_2[-2])): + skip_if_trt_rtx_turing(self, "cdist whose converter emits a GEMM") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) @@ -85,6 +97,9 @@ def forward(self, x1, x2): def test_cdist_efficiency_p_2_compute_mode( self, name, shape_1, shape_2, p, compute_mode ): + if cdist_emits_matmul(p, compute_mode, (shape_1[-2], shape_2[-2])): + skip_if_trt_rtx_turing(self, "cdist whose converter emits a GEMM") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) diff --git a/tests/py/dynamo/conversion/test_convolution_aten.py b/tests/py/dynamo/conversion/test_convolution_aten.py index 8af11cb180..fcbb2479fa 100644 --- a/tests/py/dynamo/conversion/test_convolution_aten.py +++ b/tests/py/dynamo/conversion/test_convolution_aten.py @@ -3,7 +3,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestConvolutionConverter(DispatchTestCase): @@ -233,6 +233,8 @@ def test_conv3d( groups=1, bias=True, ): + skip_if_trt_rtx_turing(self, "3D convolution") + class TestModule(torch.nn.Module): def __init__(self): super().__init__() @@ -255,6 +257,8 @@ def forward(self, x): # AssertionError: Channel dim can't be dynamic for convolution. def test_conv3d_with_dynamic_shape(self): + skip_if_trt_rtx_turing(self, "3D convolution") + class TestModule(torch.nn.Module): def __init__(self): super().__init__() diff --git a/tests/py/dynamo/conversion/test_index_put_aten.py b/tests/py/dynamo/conversion/test_index_put_aten.py index 5536c63a3c..8acde06f5d 100644 --- a/tests/py/dynamo/conversion/test_index_put_aten.py +++ b/tests/py/dynamo/conversion/test_index_put_aten.py @@ -5,7 +5,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import ENABLED_FEATURES -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing # NOTE: accumulate=True with *duplicate* indices is NOT supported in TRT. # TensorRT's ScatterMode.ND overwrites on collision — there is no scatter_add @@ -414,6 +414,12 @@ def test_index_put( if accumulate and ENABLED_FEATURES.tensorrt_rtx: pytest.skip("ScatterAdd plugin not available in TRT RTX") + # Keyed on the case's dtype rather than its index: the generated ids are bare + # test_index_put_, and there is a commented-out param above, so any index-based + # list would drift the moment a case is added or restored. + if torch.bfloat16 in (source_tensor.dtype, value_tensor.dtype): + skip_if_trt_rtx_turing(self, "bfloat16") + @torch._dynamo.assume_constant_result def get_indices_tensor(): return indices_tensor diff --git a/tests/py/dynamo/conversion/test_matmul_aten.py b/tests/py/dynamo/conversion/test_matmul_aten.py index cf1fa36e82..6826f4afff 100644 --- a/tests/py/dynamo/conversion/test_matmul_aten.py +++ b/tests/py/dynamo/conversion/test_matmul_aten.py @@ -4,7 +4,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestMatMulConverter(DispatchTestCase): @@ -28,6 +28,8 @@ class TestMatMulConverter(DispatchTestCase): ] ) def test_matmul_dot(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def __init__(self): super().__init__() @@ -83,6 +85,8 @@ def forward(self, input): ] ) def test_matmul_mm(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def __init__(self): super().__init__() @@ -123,6 +127,8 @@ def forward(self, input): ] ) def test_matmul_mv(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def __init__(self): super().__init__() @@ -150,6 +156,8 @@ def forward(self, input): ] ) def test_matmul_matmul(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def forward(self, input, other): return torch.ops.aten.matmul(input, other) @@ -181,6 +189,8 @@ def forward(self, input, other): ] ) def test_matmul_matmul_dynamic_shape(self, *args): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def forward(self, input, other): return torch.ops.aten.matmul(input, other) diff --git a/tests/py/dynamo/conversion/test_where_aten.py b/tests/py/dynamo/conversion/test_where_aten.py index 0d33536a32..d3c1dc6674 100644 --- a/tests/py/dynamo/conversion/test_where_aten.py +++ b/tests/py/dynamo/conversion/test_where_aten.py @@ -4,7 +4,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestWhereConverter(DispatchTestCase): @@ -154,6 +154,8 @@ def forward(self, condition, x, y): ] ) def test_bf16_promotion(self, x_dtype, y_dtype): + skip_if_trt_rtx_turing(self, "bfloat16") + class Where(nn.Module): def forward(self, condition, x, y): return torch.ops.aten.where.self(condition, x, y) From 633444ab65480865c57c95a93ff9bb8c39ff32ea Mon Sep 17 00:00:00 2001 From: tejaswinp Date: Sat, 29 Aug 2026 10:39:25 -0700 Subject: [PATCH 3/3] test: cover the TensorRT-RTX Turing capability guards Adds the regression tests for the fallbacks introduced alongside the Turing (SM 7.5) capability guards: FP32 GEMM, 3D convolution and bfloat16 must fall back to PyTorch, while FP16 GEMM, 2D convolution and transposed 3D convolution must stay on TensorRT. Most of the coverage is written against target_compute_capabilities=[(7, 5)], which forces Turing's partitioning on any GPU, so the guards are exercised in CI without Turing hardware. A second class repeats the same checks natively and is skipped off SM 7.5; it also pins the two failure modes that motivated the guards -- a null execution context under static shapes and a silently all-zero result under dynamic shapes. Two further cases assert on the builder config rather than on compile success, which is what makes them meaningful off Turing: everywhere except SM 7.5 an undeclared compute capability is silent. --- .../models/test_turing_capability_guards.py | 362 ++++++++++++++++++ 1 file changed, 362 insertions(+) create mode 100644 tests/py/dynamo/models/test_turing_capability_guards.py diff --git a/tests/py/dynamo/models/test_turing_capability_guards.py b/tests/py/dynamo/models/test_turing_capability_guards.py new file mode 100644 index 0000000000..e7f5044f40 --- /dev/null +++ b/tests/py/dynamo/models/test_turing_capability_guards.py @@ -0,0 +1,362 @@ +import unittest +from unittest import mock + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch_tensorrt as torchtrt +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt._features import ENABLED_FEATURES +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.conversion._TRTInterpreter import ( + TRTInterpreter, + set_rtx_compute_capabilities, +) + +TURING = (7, 5) + + +def _is_turing() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability() == TURING + + +def _trt_submodule_count(compiled) -> int: + """Number of TensorRT engine submodules in a compiled module. + + Matches on class name so the count is independent of which runtime variant + (C++, CUDA-graphs, ...) was selected. + """ + return sum( + 1 for _, m in compiled.named_modules() if "TensorRTModule" in type(m).__name__ + ) + + +class MatMul(nn.Module): + def forward(self, a, b): + return torch.matmul(a, b) + + +class Addmm(nn.Module): + """aten.addmm.default: the fused bias-add form of a GEMM, guarded on the same target.""" + + def forward(self, inp, mat1, mat2): + return torch.ops.aten.addmm.default(inp, mat1, mat2) + + +class Conv3d(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv3d(4, 8, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class ConvTranspose3d(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.ConvTranspose3d(4, 8, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class Conv2d(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(4, 8, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "Turing capability guards only apply to TensorRT-RTX", +) +class TestTuringCapabilityGuards(TestCase): + """TensorRT-RTX cannot serve some ops on Turing (SM 7.5); they must fall back. + + Per the TensorRT-RTX support matrix, Turing does not support FP32 GEMMs or 3D + convolutions, and it has no bfloat16 hardware. Without a fallback these produce a + null execution context, a segfault, or -- for dynamic-shape FP32 GEMMs -- a + silently all-zero result. + """ + + def _compile(self, mod, inputs, **kwargs): + return torchtrt.compile( + mod.eval().cuda(), + ir="dynamo", + inputs=list(inputs), + min_block_size=1, + cache_built_engines=False, + reuse_cached_engines=False, + use_python_runtime=True, + **kwargs, + ) + + def _assert_matches_eager(self, mod, compiled, inputs): + with torch.no_grad(): + ref = mod.eval().cuda()(*inputs) + out = compiled(*inputs) + self.assertFalse( + bool(torch.all(out == 0).item()) and not bool(torch.all(ref == 0).item()), + "output is all zeros while eager is not -- silent corruption", + ) + cos = F.cosine_similarity( + ref.flatten().unsqueeze(0).float(), out.flatten().unsqueeze(0).float() + ) + self.assertGreater(cos.item(), 0.99) + + # -- declared-target tests: these run on ANY GPU ------------------------------ + # Declaring Turing as a build target must produce Turing's partitioning even on + # non-Turing hardware. This is what makes the guards testable without a Turing GPU. + + def test_declared_turing_target_falls_back_fp32_gemm(self): + mod = MatMul() + inputs = (torch.randn(4, 8).cuda(), torch.randn(8, 16).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_falls_back_fp32_addmm(self): + # addmm survives lowering as its own node, so the GEMM guard sees it directly. + mod = Addmm() + inputs = ( + torch.randn(4, 6).cuda(), + torch.randn(4, 5).cuda(), + torch.randn(5, 6).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_falls_back_conv3d(self): + mod = Conv3d() + 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_fp16_gemm_on_trt(self): + mod = MatMul() + inputs = ( + torch.randn(4, 8, dtype=torch.half).cuda(), + torch.randn(8, 16, dtype=torch.half).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_conv2d_on_trt(self): + mod = Conv2d() + inputs = (torch.randn(1, 4, 16, 16).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertGreater(_trt_submodule_count(compiled), 0) + + 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() + inputs = (torch.randn(1, 4, 8, 8, 8).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertGreater(_trt_submodule_count(compiled), 0) + + @unittest.skipIf( + _is_turing(), "on Turing the native path is already guarded; see the SM75 tests" + ) + def test_non_turing_default_keeps_fp32_gemm_on_trt(self): + # Guards must not fire when Turing is not targeted. + mod = MatMul() + inputs = (torch.randn(4, 8).cuda(), torch.randn(8, 16).cuda()) + compiled = self._compile(mod, inputs) + self.assertGreater(_trt_submodule_count(compiled), 0) + + def test_target_compute_capabilities_is_engine_invariant(self): + # A cached engine built for one target set must never be reused for another. + from torch_tensorrt.dynamo._settings import _SETTINGS_TO_BE_ENGINE_INVARIANT + + self.assertIn("target_compute_capabilities", _SETTINGS_TO_BE_ENGINE_INVARIANT) + + # -- engine targeting resolves the same setting partitioning does -------------- + # Asserting on the builder config rather than on compile success is what makes these + # runnable off Turing, where an unset capability drifts silently instead of failing. + + def _compute_capability_counts(self, mod, inputs, **kwargs): + """num_compute_capabilities of every builder config this compile produced.""" + counts = [] + original = TRTInterpreter._populate_trt_builder_config + + def spy(interpreter, *args, **kwargs_): + builder_config = original(interpreter, *args, **kwargs_) + counts.append(builder_config.num_compute_capabilities) + return builder_config + + with mock.patch.object(TRTInterpreter, "_populate_trt_builder_config", spy): + self._compile(mod, inputs, **kwargs) + self.assertTrue(counts, "no engine was built, so nothing was asserted") + return counts + + def test_refittable_build_targets_a_compute_capability_by_default(self): + # Refittable builds are the ones that fail when the capability is left unset. + counts = self._compute_capability_counts( + Conv2d(), (torch.randn(1, 4, 8, 8).cuda(),), immutable_weights=False + ) + self.assertTrue(all(count == 1 for count in counts), counts) + + def test_declared_target_sets_a_capability_per_declared_target(self): + targets = [TURING] + counts = self._compute_capability_counts( + Conv2d(), + (torch.randn(1, 4, 8, 8).cuda(),), + immutable_weights=False, + target_compute_capabilities=targets, + ) + self.assertTrue(all(count == len(targets) for count in counts), counts) + + +class _FakeBuilderConfig: + """The two members set_rtx_compute_capabilities touches, and nothing else.""" + + def __init__(self, accept=True): + self.num_compute_capabilities = 0 + self.set = [] + self._accept = accept + + def set_compute_capability(self, compute_capability, idx): + self.set.append((compute_capability, idx)) + return self._accept + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx, + "set_rtx_compute_capabilities is TensorRT-RTX only", +) +class TestSetRtxComputeCapabilities(TestCase): + """Unit tests for the builder-config helper -- no GPU, no engine build.""" + + def setUp(self): + # Imported here, not at module scope: on TensorRT-RTX the ``tensorrt`` alias is + # only registered once torch_tensorrt has been imported. + import tensorrt + + self.trt = tensorrt + + def test_undeclared_targets_the_current_device(self): + config = _FakeBuilderConfig() + set_rtx_compute_capabilities(config, None) + self.assertEqual(config.num_compute_capabilities, 1) + self.assertEqual(config.set, [(self.trt.ComputeCapability.CURRENT, 0)]) + + def test_empty_list_is_treated_as_undeclared(self): + config = _FakeBuilderConfig() + set_rtx_compute_capabilities(config, []) + self.assertEqual(config.set, [(self.trt.ComputeCapability.CURRENT, 0)]) + + def test_declared_targets_are_set_by_name_in_order(self): + config = _FakeBuilderConfig() + set_rtx_compute_capabilities(config, [TURING, (8, 9)]) + self.assertEqual(config.num_compute_capabilities, 2) + self.assertEqual( + config.set, + [ + (self.trt.ComputeCapability.SM75, 0), + (self.trt.ComputeCapability.SM89, 1), + ], + ) + + def test_unknown_target_raises(self): + with self.assertRaisesRegex(ValueError, "SM99"): + set_rtx_compute_capabilities(_FakeBuilderConfig(), [(9, 9)]) + + def test_a_refused_capability_raises(self): + with self.assertRaisesRegex(RuntimeError, "SM75"): + set_rtx_compute_capabilities(_FakeBuilderConfig(accept=False), [TURING]) + + +@unittest.skipIf( + ENABLED_FEATURES.tensorrt_rtx, + "the rejection only applies to standard TensorRT", +) +class TestTargetComputeCapabilitiesRejectedOffRtx(TestCase): + def test_standard_tensorrt_rejects_declared_targets(self): + with self.assertRaisesRegex(ValueError, "only supported on TensorRT-RTX"): + CompilationSettings(target_compute_capabilities=[TURING]) + + def test_standard_tensorrt_accepts_the_default(self): + self.assertIsNone(CompilationSettings().target_compute_capabilities) + + +@unittest.skipIf( + not ENABLED_FEATURES.tensorrt_rtx or not _is_turing(), + "requires TensorRT-RTX on a Turing (SM 7.5) device", +) +class TestTuringNativeFallback(TestCase): + """Same guards, exercised natively on Turing hardware.""" + + def _compile(self, mod, inputs, **kwargs): + return torchtrt.compile( + mod.eval().cuda(), + ir="dynamo", + inputs=list(inputs), + min_block_size=1, + cache_built_engines=False, + reuse_cached_engines=False, + use_python_runtime=True, + **kwargs, + ) + + def test_fp32_gemm_static_falls_back(self): + mod = MatMul() + inputs = (torch.randn(4, 8).cuda(), torch.randn(8, 16).cuda()) + compiled = self._compile(mod, inputs) + self.assertEqual(_trt_submodule_count(compiled), 0) + + def test_fp32_gemm_dynamic_is_not_silently_zero(self): + # Unguarded, this returned an all-zero tensor of the right shape and dtype. + mod = MatMul().eval().cuda() + a = torch.randn(4, 8).cuda() + b = torch.randn(8, 16).cuda() + with torch.no_grad(): + ref = mod(a, b) + dim = torch.export.Dim("m", min=1, max=64) + ep = torch.export.export(mod, (a, b), dynamic_shapes=({0: dim}, {})) + compiled = torchtrt.compile( + ep, + arg_inputs=[a, b], + ir="dynamo", + min_block_size=1, + cache_built_engines=False, + reuse_cached_engines=False, + use_python_runtime=True, + ) + with torch.no_grad(): + out = compiled(a, b) + self.assertFalse(bool(torch.all(out == 0).item())) + cos = F.cosine_similarity( + ref.flatten().unsqueeze(0), out.flatten().unsqueeze(0) + ) + self.assertGreater(cos.item(), 0.99) + + def test_conv3d_falls_back(self): + mod = Conv3d() + inputs = (torch.randn(1, 4, 8, 8, 8).cuda(),) + compiled = self._compile(mod, inputs) + self.assertEqual(_trt_submodule_count(compiled), 0) + + def test_bfloat16_falls_back_without_crashing(self): + # Unguarded, compiling bfloat16 for SM 7.5 segfaulted the process. + class Add(nn.Module): + def forward(self, x): + return x + 1.0 + + mod = Add() + inputs = (torch.randn(4, 8, dtype=torch.bfloat16).cuda(),) + compiled = self._compile(mod, inputs, enabled_precisions={torch.bfloat16}) + self.assertEqual(_trt_submodule_count(compiled), 0) + with torch.no_grad(): + out = compiled(*inputs) + self.assertEqual(out.dtype, torch.bfloat16) + + +if __name__ == "__main__": + run_tests()