Skip to content

Add DRIVE OS support for Torch-TensorRT ExecuTorch deployment - #4735

Open
apbose wants to merge 20 commits into
mainfrom
abose/driveos-724-trt10.16-review
Open

apbose wants to merge 20 commits into
mainfrom
abose/driveos-724-trt10.16-review

Conversation

@apbose

@apbose apbose commented Sep 23, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Add an opt-in DRIVE OS / Thor build configuration for the existing Torch-TensorRT + ExecuTorch deployment flow. Development began on DRIVE OS 7.2.4; the later target validation used DRIVE OS 7.2.5. The build selects the platform SDK rather than fixing support to one DRIVE or TensorRT release.

Export models in a Python/PyTorch development environment, then load the resulting programs in a native C++ ExecuTorch application. TensorRT regions run through TensorRTBackend; supported fallback regions can run through CudaBackend / AOTInductor. This PR integrates the DRIVE build environment with these existing backends.

Python/PyTorch + Python-only Torch-TensorRT export
    -> .pte (+ external .ptd data when required)
    -> native C++ ExecuTorch application on DRIVE
         |- TensorRTBackend -> TensorRT engines
         `- CudaBackend     -> AOTInductor CUDA code

“Python-only” describes the Torch-TensorRT build used for export; PyTorch/LibTorch is still present in that authoring environment. Native C++ inference is the Python-free and LibTorch-free deployment path.

Changes

  • Add --config=driveos for Linux AArch64 dependency selection and TORCHTRT_TARGET_PLATFORM=driveos for Python packaging.
  • Keep DRIVE CUDA and TensorRT repositories separate from the normal Linux/SBSA, JetPack and Windows repositories. @cuda_driveos exposes the injected common/Thor headers and CUDA runtime directly. The Bazel DRIVE CUDA rule defaults to /usr/local/cuda and validates the DRIVE common/Thor headers, nvcc, and CUDA runtime. TORCHTRT_DRIVE_CUDA_ROOT is an optional override for a different toolkit location; TORCHTRT_DRIVE_CUDA_LIB_DIR is optional when runtime libraries are supplied separately.
  • Use the container's local TensorRT development installation under /usr for @tensorrt_driveos, following the existing NGC layout. This does not install files into /usr or download the SBSA TensorRT archive. CMake retains the standard TensorRT_ROOT override for other SDK locations.
  • Derive the DRIVE TensorRT Python requirement from the installed tensorrt distribution. The PyTorch requirement follows the current ordinary-platform policy; the runtime-wheel build uses the platform TensorRT package. The existing DLFW CI dependency-management bypass remains.
  • Extend CMake TensorRT discovery for multiarch directories and numeric or enterprise-alias version headers. Add opt-in reference-runner adjustments for the split CUDA include trees and the pinned shims' unused cuRAND host-library link.
  • Add packaging, SDK-selection and CMake discovery regression coverage. Preserve upstream's ExecuTorch CUDA allocator dependency in the DRIVE build.

SDK discovery and dependency metadata do not establish version/ABI compatibility with every DRIVE or TensorRT release. No new delegate or blanket PyTorch-operator coverage is introduced.

Validation

Historical DRIVE OS 7.2.5 target validation — October 12, 2026

These results were obtained before the latest rebase, at Torch-TensorRT commit 92e68d48e7eb1eebefc61b163afc8d91e61f717a.

Component Tested configuration
Target DRIVE Thor, DRIVE OS 7.2.5.0-46229577
Authoring container NGC PyTorch 26.06-py3
Target CUDA toolkit / NVCC CUDA 13.3 / 13.3.33
Native TensorRT library loaded DRIVE-provided 11.0.1.7
Container TensorRT Python package 11.0.0.114
PyTorch 2.16.0.dev20261009+cu132
ExecuTorch Python package 1.6.0.dev20261009+cu132
ExecuTorch native source a9665e25097b324a7116efa2bd6602287c9c064f
  • Python-only Torch-TensorRT exported TRT-only and mixed TRT+AOTInductor smoke models.
  • The CMake runner executed both programs outside Docker with target-provided TensorRT/CUDA libraries. Dependency and loader traces contained no Python, LibTorch or preserved NGC supplemental libraries. These native smoke checks compared shapes and the printed output samples.
  • The companion runtime wheel built and installed. Python TRT-only and mixed .pte execution passed full-output comparisons; these Python checks used PyTorch and are separate from the native deployment claim.
  • Additional reference examples exercised a green-context stream, staged and zero-copy KV-cache persistence, and device-resident input/output allocations.
  • The mixed smoke model was weightless. This test alone did not establish learned AOTI external-weight loading.

The authoring setup preserved supplemental NGC libraries hidden by NVIDIA runtime injection. The test also recorded a failing standalone Bazel backend analysis and package dependency conflicts. Subsequent source changes address dependency selection and metadata, but the latest source still requires a fresh complete target run. The recorded container/native TensorRT and CUDA versions are evidence for this tested configuration, not a general compatibility guarantee.

Earlier DRIVE OS 7.2.4 development used NGC 26.05, TensorRT 10.16.1.11 and CUDA 13.2. Those results remain historical; they are not evidence that the latest rebased source has been retested on 7.2.4.

Latest source and local checks — October 29, 2026

  • Review head: 2a2a856acac22330d5136d97c32c040431a13a88, rebased onto main at 39a0de8b5a145b999496f9535b6ca605e3d378c7.
  • The current upstream pin is ExecuTorch 1.6.0.dev20261021, source 5e379f610e034da507b48da61464679609517b52. This differs from the target-validation environment above.
  • Five toolchain tests and fourteen targeted packaging/CMake checks passed locally (the latter executed the actual helper/test bodies extracted from the source, without importing GPU packages). Syntax, formatting and whitespace checks were completed for the rebase changes.
  • Eight isolated Bazel 8.4.2 discovery checks passed with synthetic SDK markers: default alias, empty override, explicit root, invalid override, missing Thor headers, missing default root, missing runtime library, and a separate library directory. These validate repository discovery, not CUDA compilation or GPU execution. Ordinary Linux/SBSA CUDA selection is unchanged.
  • Fresh DRIVE GPU export, native runner and runtime-wheel validation of this head remains pending. Local checks do not substitute for target execution.

Follow-up work

  • Repeat dependency-resolving installation, the standalone Bazel backend build, CMake runner build and runtime-wheel checks on the final source snapshot.
  • Repeat TRT-only and mixed export/inference outside the authoring container with matching exporter/runtime revisions and record loaded SDK libraries.
  • Update README examples to distinguish the historical 7.2.4 setup from the later 7.2.5 validation configuration and document optional CUDA overrides.
  • Define regression coverage for subsequent DRIVE releases.

Representative model examples are separate follow-up branches: BEVFormer-Tiny and SegFormer-B0. Their model code, export fixes, numerical validation and measurements are not part of this build-support PR.

The source branch retains its historical abose/driveos-724-trt10.16-review name. Scope remains the ExecuTorch deployment path, not a complete port of every native Torch-TensorRT compiler/runtime target or validation of all DRIVE releases.

@meta-cla meta-cla Bot added the cla signed label Sep 23, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: core Issues re: The core compiler component: build system Issues re: Build system component: api [Python] Issues re: Python API component: api [C++] Issues re: C++ API component: runtime labels Sep 23, 2026
@github-actions
github-actions Bot requested a review from narendasan September 23, 2026 16:15
@lanluo-nvidia lanluo-nvidia modified the milestones: v2.15.0, v2.16.0 Sep 29, 2026
@github-actions github-actions Bot added the component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths label Oct 1, 2026
@apbose
apbose force-pushed the abose/driveos-724-trt10.16-review branch from 9674a64 to fb1ca00 Compare October 1, 2026 22:15

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-01 22:15:51.446168+00:00
+++ /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-01 22:16:28.941515+00:00
@@ -104,14 +104,11 @@
if WINDOWS_ON_ARM:
    # setuptools does not know this project-specific option, so consume it
    # before command-line processing. The environment variable remains
    # supported for CI and scripted builds.
    sys.argv.remove("--windows-on-arm")
-FORCED_WINDOWS_ARM64 = (
-    WINDOWS_ON_ARM
-    or TARGET_PLATFORM == "windows-arm64"
-)
+FORCED_WINDOWS_ARM64 = WINDOWS_ON_ARM or TARGET_PLATFORM == "windows-arm64"
if FORCED_WINDOWS_ARM64 and not IS_WINDOWS:
    raise RuntimeError("Windows ARM64 builds are supported only on Windows")
if IS_DRIVEOS and not IS_AARCH64:
    raise RuntimeError("DRIVE OS builds are supported only on Linux aarch64")
TARGET_WINDOWS_ARM64 = NATIVE_WINDOWS_ARM64 or FORCED_WINDOWS_ARM64
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-01 22:15:51.464168+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-01 22:16:34.624692+00:00
@@ -2614,11 +2614,14 @@
@pytest.mark.unit
def test_runtime_wheel_selects_driveos_bazel_config():
    source = ast.unparse(
        _class_method_def(_runtime_setup_tree(), "BazelBuild", "_build")
    )
-    assert "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'" in source
+    assert (
+        "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'"
+        in source
+    )


@pytest.mark.unit
def test_tensorrt_cmake_finder_parses_enterprise_version_macros(tmp_path):
    cmake = shutil.which("cmake")

@apbose
apbose force-pushed the abose/driveos-724-trt10.16-review branch from fb1ca00 to 96aee61 Compare October 6, 2026 02:06

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-06 02:06:41.217498+00:00
+++ /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-06 02:07:13.541907+00:00
@@ -104,14 +104,11 @@
if WINDOWS_ON_ARM:
    # setuptools does not know this project-specific option, so consume it
    # before command-line processing. The environment variable remains
    # supported for CI and scripted builds.
    sys.argv.remove("--windows-on-arm")
-FORCED_WINDOWS_ARM64 = (
-    WINDOWS_ON_ARM
-    or TARGET_PLATFORM == "windows-arm64"
-)
+FORCED_WINDOWS_ARM64 = WINDOWS_ON_ARM or TARGET_PLATFORM == "windows-arm64"
if FORCED_WINDOWS_ARM64 and not IS_WINDOWS:
    raise RuntimeError("Windows ARM64 builds are supported only on Windows")
if IS_DRIVEOS and not IS_AARCH64:
    raise RuntimeError("DRIVE OS builds are supported only on Linux aarch64")
TARGET_WINDOWS_ARM64 = NATIVE_WINDOWS_ARM64 or FORCED_WINDOWS_ARM64
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-06 02:06:41.235498+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-06 02:07:18.650488+00:00
@@ -2614,11 +2614,14 @@
@pytest.mark.unit
def test_runtime_wheel_selects_driveos_bazel_config():
    source = ast.unparse(
        _class_method_def(_runtime_setup_tree(), "BazelBuild", "_build")
    )
-    assert "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'" in source
+    assert (
+        "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'"
+        in source
+    )


@pytest.mark.unit
def test_tensorrt_cmake_finder_parses_enterprise_version_macros(tmp_path):
    cmake = shutil.which("cmake")

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-07 16:35:15.595642+00:00
+++ /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-07 16:35:35.506309+00:00
@@ -104,14 +104,11 @@
if WINDOWS_ON_ARM:
    # setuptools does not know this project-specific option, so consume it
    # before command-line processing. The environment variable remains
    # supported for CI and scripted builds.
    sys.argv.remove("--windows-on-arm")
-FORCED_WINDOWS_ARM64 = (
-    WINDOWS_ON_ARM
-    or TARGET_PLATFORM == "windows-arm64"
-)
+FORCED_WINDOWS_ARM64 = WINDOWS_ON_ARM or TARGET_PLATFORM == "windows-arm64"
if FORCED_WINDOWS_ARM64 and not IS_WINDOWS:
    raise RuntimeError("Windows ARM64 builds are supported only on Windows")
if IS_DRIVEOS and not IS_AARCH64:
    raise RuntimeError("DRIVE OS builds are supported only on Linux aarch64")
TARGET_WINDOWS_ARM64 = NATIVE_WINDOWS_ARM64 or FORCED_WINDOWS_ARM64
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-07 16:35:15.607726+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-07 16:35:38.246336+00:00
@@ -2627,11 +2627,14 @@
@pytest.mark.unit
def test_runtime_wheel_selects_driveos_bazel_config():
    source = ast.unparse(
        _class_method_def(_runtime_setup_tree(), "BazelBuild", "_build")
    )
-    assert "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'" in source
+    assert (
+        "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'"
+        in source
+    )


@pytest.mark.unit
def test_tensorrt_cmake_finder_parses_enterprise_version_macros(tmp_path):
    cmake = shutil.which("cmake")

@apbose
apbose marked this pull request as ready for review October 8, 2026 21:29

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-08 21:30:24.174476+00:00
+++ /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-08 21:30:53.445776+00:00
@@ -104,14 +104,11 @@
if WINDOWS_ON_ARM:
    # setuptools does not know this project-specific option, so consume it
    # before command-line processing. The environment variable remains
    # supported for CI and scripted builds.
    sys.argv.remove("--windows-on-arm")
-FORCED_WINDOWS_ARM64 = (
-    WINDOWS_ON_ARM
-    or TARGET_PLATFORM == "windows-arm64"
-)
+FORCED_WINDOWS_ARM64 = WINDOWS_ON_ARM or TARGET_PLATFORM == "windows-arm64"
if FORCED_WINDOWS_ARM64 and not IS_WINDOWS:
    raise RuntimeError("Windows ARM64 builds are supported only on Windows")
if IS_DRIVEOS and not IS_AARCH64:
    raise RuntimeError("DRIVE OS builds are supported only on Linux aarch64")
TARGET_WINDOWS_ARM64 = NATIVE_WINDOWS_ARM64 or FORCED_WINDOWS_ARM64
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-08 21:30:24.191856+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-08 21:30:57.951163+00:00
@@ -2627,11 +2627,14 @@
@pytest.mark.unit
def test_runtime_wheel_selects_driveos_bazel_config():
    source = ast.unparse(
        _class_method_def(_runtime_setup_tree(), "BazelBuild", "_build")
    )
-    assert "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'" in source
+    assert (
+        "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'"
+        in source
+    )


@pytest.mark.unit
def test_tensorrt_cmake_finder_parses_enterprise_version_macros(tmp_path):
    cmake = shutil.which("cmake")

@apbose
apbose force-pushed the abose/driveos-724-trt10.16-review branch from 2cb12a2 to 88c47a5 Compare October 8, 2026 21:32

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-08 21:32:53.913512+00:00
+++ /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-08 21:33:14.738364+00:00
@@ -104,14 +104,11 @@
if WINDOWS_ON_ARM:
    # setuptools does not know this project-specific option, so consume it
    # before command-line processing. The environment variable remains
    # supported for CI and scripted builds.
    sys.argv.remove("--windows-on-arm")
-FORCED_WINDOWS_ARM64 = (
-    WINDOWS_ON_ARM
-    or TARGET_PLATFORM == "windows-arm64"
-)
+FORCED_WINDOWS_ARM64 = WINDOWS_ON_ARM or TARGET_PLATFORM == "windows-arm64"
if FORCED_WINDOWS_ARM64 and not IS_WINDOWS:
    raise RuntimeError("Windows ARM64 builds are supported only on Windows")
if IS_DRIVEOS and not IS_AARCH64:
    raise RuntimeError("DRIVE OS builds are supported only on Linux aarch64")
TARGET_WINDOWS_ARM64 = NATIVE_WINDOWS_ARM64 or FORCED_WINDOWS_ARM64
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-08 21:32:53.926417+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/executorch/test_api.py	2026-10-08 21:33:17.727963+00:00
@@ -2627,11 +2627,14 @@
@pytest.mark.unit
def test_runtime_wheel_selects_driveos_bazel_config():
    source = ast.unparse(
        _class_method_def(_runtime_setup_tree(), "BazelBuild", "_build")
    )
-    assert "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'" in source
+    assert (
+        "build_config = 'driveos' if TARGET_PLATFORM == 'driveos' else 'linux'"
+        in source
+    )


@pytest.mark.unit
def test_tensorrt_cmake_finder_parses_enterprise_version_macros(tmp_path):
    cmake = shutil.which("cmake")
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/runtime/test_compile_memory.py	2026-10-08 21:32:53.934396+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/runtime/test_compile_memory.py	2026-10-08 21:33:18.805099+00:00
@@ -18,21 +18,23 @@
        env = {
            k: v
            for k, v in os.environ.items()
            if k != "TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM"
        }
-        with mock.patch.object(
-            utils, "_glibc_malloc_trim", return_value=malloc_trim
-        ), mock.patch.dict(os.environ, env, clear=True):
+        with (
+            mock.patch.object(utils, "_glibc_malloc_trim", return_value=malloc_trim),
+            mock.patch.dict(os.environ, env, clear=True),
+        ):
            utils.trim_host_heap()
        malloc_trim.assert_called_once_with(0)

    def test_env_var_disables_trimming(self):
        malloc_trim = mock.Mock(return_value=1)
-        with mock.patch.object(
-            utils, "_glibc_malloc_trim", return_value=malloc_trim
-        ), mock.patch.dict(os.environ, {"TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM": "0"}):
+        with (
+            mock.patch.object(utils, "_glibc_malloc_trim", return_value=malloc_trim),
+            mock.patch.dict(os.environ, {"TORCHTRT_ENABLE_BUILDER_MALLOC_TRIM": "0"}),
+        ):
            utils.trim_host_heap()
        malloc_trim.assert_not_called()

    def test_missing_malloc_trim_is_a_no_op(self):
        with mock.patch.object(utils, "_glibc_malloc_trim", return_value=None):

@@ -0,0 +1,31 @@
"""Expose the local DRIVE TensorRT SDK as @tensorrt_driveos, without downloads."""

def _drive_tensorrt_repository_impl(ctx):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this path change from DOS to DOS version or something? why do we need the env variable?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed the custom DRIVE TensorRT repository rule and its environment-variable override. DRIVE now uses new_local_repository(path="/usr/") with the existing local TensorRT BUILD file, following the NGC container layout. CMake retains its standard TensorRT_ROOT override for non-system SDK installations.

@apbose
apbose force-pushed the abose/driveos-724-trt10.16-review branch from 88c47a5 to 92e68d4 Compare October 12, 2026 18:50

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-12 18:51:05.466751+00:00
+++ /home/runner/work/TensorRT/TensorRT/setup.py	2026-10-12 18:51:38.741731+00:00
@@ -104,14 +104,11 @@
if WINDOWS_ON_ARM:
    # setuptools does not know this project-specific option, so consume it
    # before command-line processing. The environment variable remains
    # supported for CI and scripted builds.
    sys.argv.remove("--windows-on-arm")
-FORCED_WINDOWS_ARM64 = (
-    WINDOWS_ON_ARM
-    or TARGET_PLATFORM == "windows-arm64"
-)
+FORCED_WINDOWS_ARM64 = WINDOWS_ON_ARM or TARGET_PLATFORM == "windows-arm64"
if FORCED_WINDOWS_ARM64 and not IS_WINDOWS:
    raise RuntimeError("Windows ARM64 builds are supported only on Windows")
if IS_DRIVEOS and not IS_AARCH64:
    raise RuntimeError("DRIVE OS builds are supported only on Linux aarch64")
TARGET_WINDOWS_ARM64 = NATIVE_WINDOWS_ARM64 or FORCED_WINDOWS_ARM64

@apbose
apbose force-pushed the abose/driveos-724-trt10.16-review branch from 92e68d4 to ecde4ea Compare October 29, 2026 20:30
@apbose apbose changed the title Add DRIVE OS 7.2.4 build support for ExecuTorch deployment Add DRIVE OS support for Torch-TensorRT ExecuTorch deployment Oct 29, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [C++] Issues re: C++ API component: api [Python] Issues re: Python API component: build system Issues re: Build system component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: runtime component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants