fix: fall back to PyTorch for ops TensorRT-RTX cannot serve on Turing (SM 7.5)

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,49 @@ 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 -------------------------------------
+#
+# TensorRT-RTX runs on SM 7.5 and up, but its support matrix carves Turing out of
+# several paths: FP32 GEMMs and 3D convolutions are documented as unsupported on
+# compute capability 7.5, and bfloat16 has no Turing hardware at all. It also notes
+# that Turing "compatibility is not included by default because including it may
+# impact model performance", so targeting Turing is opt-in.
+#
+# These helpers exist so capability validators key off the compute capabilities the
+# engine is being *built for*, not the machine doing the building. Calling
+# torch.cuda.get_device_capability() inside a validator bakes the build host into the
+# artifact, which is wrong for ahead-of-time deployment: a module compiled on Ampere
+# and shipped to Turing would retain ops Turing cannot execute.
+
+TURING_COMPUTE_CAPABILITY = (7, 5)
+
+
+def get_target_compute_capabilities(
+    settings: Optional[Any] = None,
+) -> Tuple[Tuple[int, int], ...]:
+    """Compute capabilities this compilation targets.
+
+    Returns the explicitly declared targets when the caller set them, otherwise the
+    capability of the current device.
+    """
+    targets = (
+        getattr(settings, "target_compute_capabilities", None) if settings else None
+    )
+    if targets:
+        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)

py/torch_tensorrt/dynamo/_compiler.py

@@ -98,6 +98,9 @@ def cross_compile_for_windows(
     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,
@@ -181,6 +184,7 @@ def cross_compile_for_windows(
         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
@@ -352,6 +356,7 @@ def cross_compile_for_windows(
         "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,
@@ -462,6 +467,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,
@@ -560,6 +568,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
@@ -763,6 +772,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,
@@ -1758,6 +1768,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,
@@ -1854,6 +1867,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
@@ -2033,6 +2047,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,

py/torch_tensorrt/dynamo/_defaults.py

@@ -63,6 +63,10 @@
 ENABLE_CROSS_COMPILE_FOR_WINDOWS = False
 TILING_OPTIMIZATION_LEVEL = "none"
 L2_LIMIT_FOR_TILING = -1
+# None means "target the current device". Set explicitly to build an engine deployable
+# on other architectures; see torch_tensorrt._utils for why capability validators must
+# consult this rather than the build host's device.
+TARGET_COMPUTE_CAPABILITIES = None
 USE_DISTRIBUTED_MODE_TRACE = False
 OFFLOAD_MODULE_TO_CPU = False
 ENABLE_AUTOCAST = False

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,10 @@ def _normalize_disabled_constant_fold_exclusions(
         validate_disabled_constant_fold_exclusions,
     )
 
-    return validate_disabled_constant_fold_exclusions(rule_ids)
+    # Bind to a typed local: the imported helper is untyped from mypy's view here,
+    # and returning it directly trips --strict's no-any-return.
+    normalized: Set[str] = validate_disabled_constant_fold_exclusions(rule_ids)
+    return normalized
 
 
 @dataclass
@@ -106,6 +110,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. This 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 +186,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
@@ -245,6 +253,10 @@ def __setstate__(self, state: dict[str, Any]) -> None:
     "sparse_weights",
     "engine_capability",
     "hardware_compatible",
+    # An engine built for one set of compute capabilities cannot be reused for a
+    # compile targeting a different set -- doing so would silently reintroduce ops the
+    # target architecture cannot execute.
+    "target_compute_capabilities",
     "refit_identical_engine_weights",
     "immutable_weights",
     "enable_weight_streaming",

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]] = (
@@ -385,6 +384,35 @@ def _populate_trt_builder_config(
                     self.compilation_settings.l2_limit_for_tiling
                 )
 
+        # TensorRT-RTX ahead-of-time targeting. Left unset, TensorRT-RTX compiles for
+        # whatever device is present at build time, which is right for
+        # compile-here-run-here but cannot produce an artifact for another
+        # architecture. Turing in particular is opt-in: including it by default may
+        # cost performance elsewhere. The same setting drives the capability
+        # validators, so partitioning and engine targeting cannot drift apart.
+        if (
+            ENABLED_FEATURES.tensorrt_rtx
+            and self.compilation_settings.target_compute_capabilities
+        ):
+            targets = self.compilation_settings.target_compute_capabilities
+            builder_config.num_compute_capabilities = len(targets)
+            for idx, (major, minor) in enumerate(targets):
+                name = f"SM{major}{minor}"
+                compute_capability = getattr(trt.ComputeCapability, name, None)
+                if compute_capability 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}"
+                    )
+                if not builder_config.set_compute_capability(compute_capability, idx):
+                    raise RuntimeError(
+                        f"Failed to set TensorRT-RTX compute capability {name}"
+                    )
+            _LOGGER.info(f"Targeting TensorRT-RTX compute capabilities {targets}")
+
         return builder_config
 
     def _create_timing_cache(