fix: TensorRT-RTX Turing (SM 7.5) capability guards and compute-capability targeting

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)

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,

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

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",

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<major><minor> 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(