Fix the ExecuTorch CUDA 12 build and test failures

core/runtime/TRTEngine.cpp

@@ -6,6 +6,7 @@
 #include "NvInfer.h"
 #include "c10/cuda/CUDACachingAllocator.h"
 #include "c10/cuda/CUDAStream.h"
+#include "c10/util/Type.h"
 #include "torch/csrc/jit/frontend/function_schema_parser.h"
 #include "torch/cuda.h"
 
@@ -913,7 +914,13 @@ bool TRTEngine::bind_nccl_comm() {
   TORCHTRT_CHECK(backend != nullptr, "ProcessGroup '" << this->group_name << "' has no NCCL backend");
 
   auto* nccl_pg = dynamic_cast<c10d::ProcessGroupNCCL*>(backend.get());
-  TORCHTRT_CHECK(nccl_pg != nullptr, "Backend is not ProcessGroupNCCL");
+  // Name the type that arrived. getBackend returned non-null, so something is there; without
+  // printing it the failure says only that the cast missed, which is not enough to tell a wrapper
+  // apart from a different concrete group.
+  TORCHTRT_CHECK(
+      nccl_pg != nullptr,
+      "Backend for ProcessGroup '" << this->group_name << "' is not ProcessGroupNCCL, got "
+                                   << c10::demangle(typeid(*backend).name()));
 
   at::cuda::set_device(this->device_info.id);
 

py/torch-tensorrt-executorch-runtime/setup.py

@@ -21,7 +21,6 @@
 REPO_ROOT = HERE.parents[1]
 BAZEL_TARGET = "//py/torch-tensorrt-executorch-runtime/native:delegate_native"
 BUILD_NONCE = os.getenv("TORCH_TENSORRT_EXECUTORCH_BUILD_NONCE", uuid.uuid4().hex)
-CUDA_RUNTIME_DISTRIBUTION = "nvidia-cuda-runtime"
 
 
 def torchtrt_version() -> str:
@@ -50,17 +49,41 @@ def installed_version(distribution: str) -> str:
 
 
 def tensorrt_distribution() -> str:
-    """Return the TensorRT distribution matching the PyTorch CUDA build."""
+    """Return the TensorRT distribution matching the PyTorch CUDA build.
+
+    Matched on the major. test-infra's matrix carries cu128 alongside cu126, and a 12.x that is
+    not exactly 12.6 used to reach the "Unsupported CUDA version" path here even though
+    tensorrt-cu12 is what it needs.
+    """
     cuda_version = torch.version.cuda
     if cuda_version is None:
         raise RuntimeError("CUDA-enabled PyTorch is required to build this wheel")
-    if cuda_version.startswith("12.6"):
+    if cuda_version.startswith("12."):
         return "tensorrt-cu12"
     if cuda_version.startswith("13."):
         return "tensorrt-cu13"
     raise RuntimeError(f"Unsupported CUDA version: {cuda_version}")
 
 
+def cuda_runtime_distribution() -> str:
+    """Return the CUDA runtime distribution matching the PyTorch CUDA build.
+
+    NVIDIA splits this one by major: the CUDA 12 wheels are published as
+    ``nvidia-cuda-runtime-cu12``, while the unsuffixed ``nvidia-cuda-runtime`` is the CUDA 13
+    line. Naming the unsuffixed one unconditionally made every CUDA 12 row fail with "No package
+    metadata was found for nvidia-cuda-runtime", because what torch installed there was the
+    suffixed distribution.
+    """
+    cuda_version = torch.version.cuda
+    if cuda_version is None:
+        raise RuntimeError("CUDA-enabled PyTorch is required to build this wheel")
+    if cuda_version.startswith("12."):
+        return "nvidia-cuda-runtime-cu12"
+    if cuda_version.startswith("13."):
+        return "nvidia-cuda-runtime"
+    raise RuntimeError(f"Unsupported CUDA version: {cuda_version}")
+
+
 class BazelExtension(Extension):
     def __init__(self, name: str) -> None:
         super().__init__(name, sources=[])
@@ -141,6 +164,7 @@ def build_extension(self, ext: Extension) -> None:
 
 
 TENSORRT_DISTRIBUTION = tensorrt_distribution()
+CUDA_RUNTIME_DISTRIBUTION = cuda_runtime_distribution()
 executorch_version = installed_version("executorch")
 tensorrt_version = installed_version(TENSORRT_DISTRIBUTION)
 cuda_runtime_version = installed_version(CUDA_RUNTIME_DISTRIBUTION)

tests/py/dynamo/executorch/test_api.py

@@ -231,14 +231,72 @@ def test_runtime_wheel_uses_public_torch_version():
 
 
 @pytest.mark.unit
-def test_runtime_wheel_pins_cuda_13_native_dependencies():
+def test_runtime_wheel_pins_its_native_dependencies_per_cuda_major():
+    """Both native dependencies are resolved from the CUDA the build actually uses.
+
+    Hardcoding the CUDA 13 spellings was right while only CUDA 13 shipped and wrong once 12.6 came
+    back. NVIDIA splits the runtime by major: CUDA 12 publishes nvidia-cuda-runtime-cu12 while the
+    unsuffixed name is the CUDA 13 line, so naming the unsuffixed one unconditionally failed every
+    CUDA 12 row with "No package metadata was found for nvidia-cuda-runtime". TensorRT splits the
+    same way.
+    """
     setup_source = _RUNTIME_SETUP_PY.read_text(encoding="utf-8")
-    assert 'TENSORRT_DISTRIBUTION = "tensorrt-cu13"' in setup_source
-    assert 'CUDA_RUNTIME_DISTRIBUTION = "nvidia-cuda-runtime"' in setup_source
+    assert "TENSORRT_DISTRIBUTION = tensorrt_distribution()" in setup_source, (
+        "the TensorRT distribution is no longer resolved from the build's CUDA version, so a CUDA "
+        "12 row would declare the CUDA 13 distribution"
+    )
+    assert "CUDA_RUNTIME_DISTRIBUTION = cuda_runtime_distribution()" in setup_source, (
+        "the CUDA runtime distribution is no longer resolved from the build's CUDA version, so a "
+        "CUDA 12 row would look for a distribution torch did not install"
+    )
+
+    # Execute both resolvers rather than grepping for the names. Checking that each spelling
+    # appears somewhere in the file still passes when the two mappings are swapped, which would
+    # send every CUDA 12 row after CUDA 13 packages and vice versa.
+    tree = ast.parse(setup_source)
+    wanted = ("tensorrt_distribution", "cuda_runtime_distribution")
+    functions = [
+        node
+        for node in tree.body
+        if isinstance(node, ast.FunctionDef) and node.name in wanted
+    ]
+    assert {node.name for node in functions} == set(wanted), (
+        f"expected both resolvers to be module-level functions, found "
+        f"{sorted(node.name for node in functions)}"
+    )
+    fake_torch = types.SimpleNamespace(version=types.SimpleNamespace(cuda=None))
+    namespace: dict = {"torch": fake_torch}
+    exec(
+        compile(ast.Module(body=functions, type_ignores=[]), "<setup.py>", "exec"),
+        namespace,
+    )
+
+    for cuda_version, tensorrt, cuda_runtime in (
+        ("12.6", "tensorrt-cu12", "nvidia-cuda-runtime-cu12"),
+        ("12.8", "tensorrt-cu12", "nvidia-cuda-runtime-cu12"),
+        ("13.0", "tensorrt-cu13", "nvidia-cuda-runtime"),
+        ("13.2", "tensorrt-cu13", "nvidia-cuda-runtime"),
+    ):
+        fake_torch.version.cuda = cuda_version
+        assert namespace["tensorrt_distribution"]() == tensorrt, (
+            f"CUDA {cuda_version} resolves the wrong TensorRT distribution, so that row installs "
+            "the other CUDA major's native libraries"
+        )
+        assert namespace["cuda_runtime_distribution"]() == cuda_runtime, (
+            f"CUDA {cuda_version} resolves the wrong CUDA runtime distribution, so that row looks "
+            "for metadata torch did not install"
+        )
+
+    # A CUDA-less torch and an unsupported major are refused rather than guessed at.
+    for cuda_version in (None, "11.8"):
+        fake_torch.version.cuda = cuda_version
+        for name in wanted:
+            with pytest.raises(RuntimeError):
+                namespace[name]()
+
     assert "torch=={public_version(torch.__version__)}" in setup_source
     assert "{TENSORRT_DISTRIBUTION}=={tensorrt_version}" in setup_source
     assert "{CUDA_RUNTIME_DISTRIBUTION}=={cuda_runtime_version}" in setup_source
-    assert "nvidia-cuda-runtime-cu12" not in setup_source
 
 
 @pytest.mark.unit