MD-TRT Support, Compile/Export, C++ and Python by narendasan · Pull Request #4183 · pytorch/TensorRT

narendasan · 2026-04-12T19:08:47Z

Description

Opening this to test the CI

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

- C++ runtime: NCCL communicator init via c10d, rank/world_size serialization, DynamicOutputAllocator, ABI version bump to 8 - Python runtime: distributed support in PythonTorchTensorRTModule and TorchTensorRTModule, NCCL library auto-detection - Conversion: native TRT DistCollective API (AllGather, ReduceScatter, AllReduce) with TRT-LLM plugin fallback - Graph lowering: fuse c10d_functional collectives + wait_tensor into single ops - Feature detection: native_trt_collectives flag, platform validation, graceful fallback chain - Build: conditional NCCL compilation via torch_nccl toolchain - Examples: tensor_parallel_simple_example.py, tensor_parallel_llama_llm.py

…, ABI v9

…g and enable DTensor decomposition

…hapes Five interconnected fixes: 1. fold_get_attr_item_calls: fold scalar param .item() calls into Python scalars before AOT tracing. Inside FakeTensorMode, even real-tensor .item() calls raise DataDependentOutputException. 2. backends.py: three changes: - call fold_get_attr_item_calls before entering FakeTensorMode - detect vmap/higher-order ops and route them through aot_autograd instead of aot_export_joint_simple (which doesn't handle HOPs) - on TRT build failure, strip TRT-only kwargs (use_fp32_acc) from the fallback graph before returning it to PyTorch 3. _decompositions.py: prevent SDPA from leaking back into the decomp table via Core ATen Interchange ops even after being removed from TORCH_TRT_DECOMPOSITIONS. 4. partitioning/common.py: lower the default max dynamic shape from min*2^16 to min*2^12 — 65536 is too large for TRT to find kernel implementations for attention ops. 5. _TorchTensorRTModule.py: move CPU scalar inputs to CUDA before execution — aot_autograd lifts scalar attributes (e.g. head_dim^-0.5) as explicit graph inputs; TRT requires all inputs on CUDA. Also fixes remove_sym_nodes to match tensor sources by equality rather than local_name so that GetItemSource bases (from torch.compile dynamic=True) are matched correctly, and updates register_sdpa.py to handle aten.scaled_dot_product_attention.default (the form produced after aot_autograd) in addition to the flash/efficient variants.

…a 3.2 1B

apbose · 2026-04-16T07:51:18Z

+        if id(engine) not in seen:
+            seen.add(id(engine))
+            if getattr(engine, "is_md", False):
+                engine.set_group_name(group_name)


why is this added?

We need to track instances for clean up

…stributed context usage in some tests

github-actions

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

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-18 20:14:39.307048+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-18 20:15:03.544385+00:00
@@ -386,11 +386,10 @@
                logger.debug(
                    "Barrier after execution context creation (distributed NCCL engine)"
                )
                dist.barrier()

-
        if ENABLED_FEATURES.tensorrt_rtx:
            self._setup_runtime_config()

        self.context = self._create_context()
        assert self.context is not None, "Failed to create execution context"

github-actions

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

diff --git a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.h b/tmp/changes.txt
index cd8af65..615600d 100644
--- a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.h
+++ b/tmp/changes.txt
@@ -33,17 +33,17 @@ namespace core {
namespace runtime {

using FlattenedState = std::tuple<
-    std::tuple<std::string, std::string>,  // ABI_VERSION
-    std::tuple<std::string, std::string>,  // name
-    std::tuple<std::string, std::string>,  // device
-    std::tuple<std::string, std::string>,  // engine
-    std::tuple<std::string, std::string>,  // input binding names
-    std::tuple<std::string, std::string>,  // output binding names
-    std::tuple<std::string, std::string>,  // HW compatibility
-    std::tuple<std::string, std::string>,  // requires_output_allocator
-    std::tuple<std::string, std::string>,  // serialized metadata
-    std::tuple<std::string, std::string>,  // Platform
-    std::tuple<std::string, std::string>,  // Resource Allocation Strategy
+    std::tuple<std::string, std::string>, // ABI_VERSION
+    std::tuple<std::string, std::string>, // name
+    std::tuple<std::string, std::string>, // device
+    std::tuple<std::string, std::string>, // engine
+    std::tuple<std::string, std::string>, // input binding names
+    std::tuple<std::string, std::string>, // output binding names
+    std::tuple<std::string, std::string>, // HW compatibility
+    std::tuple<std::string, std::string>, // requires_output_allocator
+    std::tuple<std::string, std::string>, // serialized metadata
+    std::tuple<std::string, std::string>, // Platform
+    std::tuple<std::string, std::string>, // Resource Allocation Strategy
    std::tuple<std::string, std::string>>; // requires_multidevice

struct TorchTRTRuntimeStates {
ERROR: Some files do not conform to style guidelines

github-actions

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

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-20 22:42:10.394317+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-20 22:42:31.904975+00:00
@@ -386,11 +386,10 @@
                logger.debug(
                    "Barrier after execution context creation (distributed NCCL engine)"
                )
                dist.barrier()

-
        if ENABLED_FEATURES.tensorrt_rtx:
            self._setup_runtime_config()

        self.context = self._create_context()
        assert self.context is not None, "Failed to create execution context"

github-actions

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

diff --git a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.h b/tmp/changes.txt
index cd8af65..615600d 100644
--- a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.h
+++ b/tmp/changes.txt
@@ -33,17 +33,17 @@ namespace core {
namespace runtime {

using FlattenedState = std::tuple<
-    std::tuple<std::string, std::string>,  // ABI_VERSION
-    std::tuple<std::string, std::string>,  // name
-    std::tuple<std::string, std::string>,  // device
-    std::tuple<std::string, std::string>,  // engine
-    std::tuple<std::string, std::string>,  // input binding names
-    std::tuple<std::string, std::string>,  // output binding names
-    std::tuple<std::string, std::string>,  // HW compatibility
-    std::tuple<std::string, std::string>,  // requires_output_allocator
-    std::tuple<std::string, std::string>,  // serialized metadata
-    std::tuple<std::string, std::string>,  // Platform
-    std::tuple<std::string, std::string>,  // Resource Allocation Strategy
+    std::tuple<std::string, std::string>, // ABI_VERSION
+    std::tuple<std::string, std::string>, // name
+    std::tuple<std::string, std::string>, // device
+    std::tuple<std::string, std::string>, // engine
+    std::tuple<std::string, std::string>, // input binding names
+    std::tuple<std::string, std::string>, // output binding names
+    std::tuple<std::string, std::string>, // HW compatibility
+    std::tuple<std::string, std::string>, // requires_output_allocator
+    std::tuple<std::string, std::string>, // serialized metadata
+    std::tuple<std::string, std::string>, // Platform
+    std::tuple<std::string, std::string>, // Resource Allocation Strategy
    std::tuple<std::string, std::string>>; // requires_multidevice

struct TorchTRTRuntimeStates {
ERROR: Some files do not conform to style guidelines

github-actions

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

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-21 00:20:02.596692+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-21 00:20:22.271000+00:00
@@ -386,11 +386,10 @@
                logger.debug(
                    "Barrier after execution context creation (distributed NCCL engine)"
                )
                dist.barrier()

-
        if ENABLED_FEATURES.tensorrt_rtx:
            self._setup_runtime_config()

        self.context = self._create_context()
        assert self.context is not None, "Failed to create execution context"

…us option

github-actions

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

diff --git a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.cpp b/tmp/changes.txt
index 4b91415..ae5232b 100644
--- a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.cpp
+++ b/tmp/changes.txt
@@ -573,14 +573,16 @@ bool TRTEngine::bind_nccl_comm() {
    } else if (nccl_groups.size() > 1) {
      std::string names;
      for (const auto& n : nccl_groups) {
-        if (!names.empty()) names += ", ";
+        if (!names.empty())
+          names += ", ";
        names += "'" + n + "'";
      }
      LOG_WARNING(
          "This TRT engine requires NCCL but multiple NCCL process groups are registered ("
-          << names << "). Cannot auto-select a group — NCCL bind deferred. "
-          "Use the recommended workflow: "
-          "with torch_tensorrt.distributed.distributed_context(group, model) as m: m(inp)");
+          << names
+          << "). Cannot auto-select a group — NCCL bind deferred. "
+             "Use the recommended workflow: "
+             "with torch_tensorrt.distributed.distributed_context(group, model) as m: m(inp)");
    } else {
      LOG_WARNING(
          "This TRT engine requires NCCL (requires_native_multidevice=true) but no NCCL process group "
ERROR: Some files do not conform to style guidelines

github-actions

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

diff --git a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.cpp b/tmp/changes.txt
index 4b91415..ae5232b 100644
--- a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.cpp
+++ b/tmp/changes.txt
@@ -573,14 +573,16 @@ bool TRTEngine::bind_nccl_comm() {
    } else if (nccl_groups.size() > 1) {
      std::string names;
      for (const auto& n : nccl_groups) {
-        if (!names.empty()) names += ", ";
+        if (!names.empty())
+          names += ", ";
        names += "'" + n + "'";
      }
      LOG_WARNING(
          "This TRT engine requires NCCL but multiple NCCL process groups are registered ("
-          << names << "). Cannot auto-select a group — NCCL bind deferred. "
-          "Use the recommended workflow: "
-          "with torch_tensorrt.distributed.distributed_context(group, model) as m: m(inp)");
+          << names
+          << "). Cannot auto-select a group — NCCL bind deferred. "
+             "Use the recommended workflow: "
+             "with torch_tensorrt.distributed.distributed_context(group, model) as m: m(inp)");
    } else {
      LOG_WARNING(
          "This TRT engine requires NCCL (requires_native_multidevice=true) but no NCCL process group "
ERROR: Some files do not conform to style guidelines

github-actions

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

diff --git a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.cpp b/tmp/changes.txt
index 4b91415..ae5232b 100644
--- a/home/runner/work/TensorRT/TensorRT/core/runtime/TRTEngine.cpp
+++ b/tmp/changes.txt
@@ -573,14 +573,16 @@ bool TRTEngine::bind_nccl_comm() {
    } else if (nccl_groups.size() > 1) {
      std::string names;
      for (const auto& n : nccl_groups) {
-        if (!names.empty()) names += ", ";
+        if (!names.empty())
+          names += ", ";
        names += "'" + n + "'";
      }
      LOG_WARNING(
          "This TRT engine requires NCCL but multiple NCCL process groups are registered ("
-          << names << "). Cannot auto-select a group — NCCL bind deferred. "
-          "Use the recommended workflow: "
-          "with torch_tensorrt.distributed.distributed_context(group, model) as m: m(inp)");
+          << names
+          << "). Cannot auto-select a group — NCCL bind deferred. "
+             "Use the recommended workflow: "
+             "with torch_tensorrt.distributed.distributed_context(group, model) as m: m(inp)");
    } else {
      LOG_WARNING(
          "This TRT engine requires NCCL (requires_native_multidevice=true) but no NCCL process group "
ERROR: Some files do not conform to style guidelines

apbose and others added 11 commits April 12, 2026 11:41

removing the try-except block in TRTengine.cpp and correcting the typis

aaa6557

Redesign distributed inference API: auto-detect rank, lazy NCCL setup…

4c1e68d

…, ABI v9

remove nccl.h dependancy

7cfa40b

clean up import and add comment

ac96255

moving setup_nccl_library call to example script

fe1c6f4

work on the save/load export part-add is_md flag, guard export tracin…

b658c7a

…g and enable DTensor decomposition

refactor: Adjusting how we use NCCL

a35dfe6

test: add torch.compile(backend='tensorrt') integration test for Llam…

2aa8f14

…a 3.2 1B

feat: llama3.2 working with MD-TRT

6f81a66

meta-cla bot added the cla signed label Apr 12, 2026

github-actions bot requested a review from zewenli98 April 12, 2026 19:09

This comment was marked as outdated.

Sign in to view

narendasan force-pushed the push-vqqzkszwrvyx branch from 67134da to b5b1f5f Compare April 12, 2026 19:15

This comment was marked as outdated.

Sign in to view

narendasan force-pushed the push-vqqzkszwrvyx branch from b5b1f5f to 1957cc4 Compare April 12, 2026 19:18

apbose reviewed Apr 16, 2026

View reviewed changes

Comment thread py/torch_tensorrt/distributed/_distributed.py

apbose reviewed Apr 16, 2026

View reviewed changes

Comment thread py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

replacing torchrun with torchtrtrun for right .so

edf6518

This comment was marked as outdated.

Sign in to view

chore: apply linting

9e390eb

narendasan force-pushed the push-vqqzkszwrvyx branch from 93277a4 to 9e390eb Compare April 16, 2026 16:57

use correct group for dummy all_reduce

1b4e559

apbose force-pushed the push-vqqzkszwrvyx branch from 04b606e to 1b4e559 Compare April 16, 2026 22:51

apbose and others added 2 commits April 16, 2026 19:33

Broaden NCCL skip guards to include native TRT collectives and fix di…

df51acf

…stributed context usage in some tests

Merge branch 'main' into push-vqqzkszwrvyx

4665692

github-actions bot requested changes Apr 18, 2026

View reviewed changes

narendasan added 3 commits April 20, 2026 15:15

fix: Update the engine cache to be aware of the new setting

fd32b5b

update TRT version

7dce75e

fix: the md property was not properly handled across the library

3d3c7d1

github-actions bot requested changes Apr 20, 2026

View reviewed changes

github-actions bot requested changes Apr 21, 2026

View reviewed changes

chore: skip test which is not valid ATen attn

fe59779

narendasan force-pushed the push-vqqzkszwrvyx branch from 9d95aa9 to fe59779 Compare April 21, 2026 00:26

narendasan added 2 commits April 20, 2026 18:38

finalizing internal apis

fdfd45a

fix: align apis, make sure to defer binding unless there is one obvio…

f9bd6a4

…us option

github-actions bot requested changes Apr 21, 2026

View reviewed changes

test: make bert test backwards compatible

34ccf1d

github-actions bot requested changes Apr 21, 2026

View reviewed changes

fix: address some issues in the converters

6379b06

github-actions bot requested changes Apr 21, 2026

View reviewed changes

Conversation

narendasan commented Apr 12, 2026

Description

Type of change

Checklist:

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

apbose Apr 16, 2026

Choose a reason for hiding this comment

Uh oh!

narendasan Apr 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants