MD-TRT Support, Compile/Export, C++ and Python

[narendasan]

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

[narendasan]

I would like to remove this but didnt have time to check if the device operations in python suppress this correctly

[narendasan]

Not entirely sure this is necessary

[narendasan]

We should only do this if there is one available group. If there are multiple NCCL groups available we should tell the user to manually select

[narendasan]

Lets dig into this more after the PR lands

[narendasan]

Move into its own file

[apbose]

this should be set before the self.setup_engine()

[apbose]

something like this works
engine_json_lower = engine_json.lower()
self._has_nccl_ops = "dist_collective" in engine_json_lower or "nccl" in engine_json_lower or "allreduce" in engine_json_lower

[github-actions]

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

--- /home/runner/work/TensorRT/TensorRT/examples/distributed_inference/tensor_parallel_simple_example_md.py	2026-04-14 20:34:53.887235+00:00
+++ /home/runner/work/TensorRT/TensorRT/examples/distributed_inference/tensor_parallel_simple_example_md.py	2026-04-14 20:35:10.238612+00:00
@@ -97,13 +97,13 @@
        x = self.out_proj2(self.relu(self.in_proj2(x)))
        return x


def get_model(device_mesh):
-    assert world_size % 2 == 0, (
-        f"TP examples require an even number of GPUs, got {world_size}"
-    )
+    assert (
+        world_size % 2 == 0
+    ), f"TP examples require an even number of GPUs, got {world_size}"
    model = ToyModel().to(DEVICE)
    parallelize_module(
        module=model,
        device_mesh=device_mesh,
        parallelize_plan={
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py	2026-04-14 20:34:53.901309+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py	2026-04-14 20:35:11.188933+00:00
@@ -22,11 +22,13 @@

if ENABLED_FEATURES.native_trt_collectives:
    # Use native TensorRT DistCollective API (no TensorRT-LLM dependency)
    _LOGGER.info("Using native TensorRT DistCollective API for distributed operations")

-    @dynamo_tensorrt_converter(tensorrt_fused_nccl_all_gather_op, requires_multidevice=True)
+    @dynamo_tensorrt_converter(
+        tensorrt_fused_nccl_all_gather_op, requires_multidevice=True
+    )
    def fused_nccl_gather(
        ctx: ConversionContext,
        target: Target,
        args: Tuple[Argument, ...],
        kwargs: Dict[str, Argument],
@@ -39,11 +41,13 @@
            SourceIR.ATEN,
            name,
            [args[0]],
        )

-    @dynamo_tensorrt_converter(tensorrt_fused_nccl_reduce_scatter_op, requires_multidevice=True)
+    @dynamo_tensorrt_converter(
+        tensorrt_fused_nccl_reduce_scatter_op, requires_multidevice=True
+    )
    def fused_nccl_reduce_scatter(
        ctx: ConversionContext,
        target: Target,
        args: Tuple[Argument, ...],
        kwargs: Dict[str, Argument],
@@ -56,11 +60,13 @@
            SourceIR.ATEN,
            name,
            [args[0]],
        )

-    @dynamo_tensorrt_converter(tensorrt_fused_nccl_all_reduce_op, requires_multidevice=True)
+    @dynamo_tensorrt_converter(
+        tensorrt_fused_nccl_all_reduce_op, requires_multidevice=True
+    )
    def fused_nccl_all_reduce(
        ctx: ConversionContext,
        target: Target,
        args: Tuple[Argument, ...],
        kwargs: Dict[str, Argument],
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py	2026-04-14 20:34:53.900894+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py	2026-04-14 20:35:11.722086+00:00
@@ -669,11 +669,11 @@
            cuda_engine,
            self._input_names,
            self._output_names,
            self.weight_name_map,
            self.ctx.requires_output_allocator,
-            self.ctx.requires_multidevice
+            self.ctx.requires_multidevice,
        )

    def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
        self._cur_node_name = get_node_name(n)
        self._cur_node = n
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-14 20:34:53.907339+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py	2026-04-14 20:35:13.284044+00:00
@@ -368,12 +368,18 @@
            # For engines with native NCCL collective layers, all ranks must
            # have a live IExecutionContext before any rank executes a
            # collective. Barrier here so a fast-compiling rank does not race
            # ahead and issue an NCCL op while another rank is still inside
            # deserialize_cuda_engine / create_execution_context.
-            if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1:
-                logger.debug("Barrier after execution context creation (distributed NCCL engine)")
+            if (
+                dist.is_available()
+                and dist.is_initialized()
+                and dist.get_world_size() > 1
+            ):
+                logger.debug(
+                    "Barrier after execution context creation (distributed NCCL engine)"
+                )
                dist.barrier()

        assert self.context is not None, "Failed to create execution context"
        assert self.engine.num_io_tensors == (
            len(self.input_names) + len(self.output_names)
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/distributed/test_native_nccl.py	2026-04-14 20:34:53.931747+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/distributed/test_native_nccl.py	2026-04-14 20:35:15.878751+00:00
@@ -1525,13 +1525,11 @@

        expected = torch.full((1, 4), float(expected_sum), device=device)
        _check_close(out, expected, f"context_switch sg{i+1} rank={rank}")


-def _multirank_pg_migration(
-    rank: int, world_size: int, device: torch.device
-) -> None:
+def _multirank_pg_migration(rank: int, world_size: int, device: torch.device) -> None:
    """Compile with the default world group, run inference, then migrate to a new
    subgroup via distributed_group(new_group, model) and verify that inference
    still produces correct results — i.e. the NCCL communicator is re-bound.

    Tests both the C++ runtime (set_group_name resets nccl_initialized) and the
@@ -1562,13 +1560,11 @@
        def __init__(self, pg_name: str) -> None:
            super().__init__()
            self.pg_name = pg_name

        def forward(self, x: torch.Tensor) -> torch.Tensor:
-            out = torch.ops._c10d_functional.all_reduce.default(
-                x, "sum", self.pg_name
-            )
+            out = torch.ops._c10d_functional.all_reduce.default(x, "sum", self.pg_name)
            return torch.ops._c10d_functional.wait_tensor.default(out)

    inp = torch.full((1, 4), float(rank + 1), device=device)
    expected_sum = world_size * (world_size + 1) // 2
    expected = torch.full((1, 4), float(expected_sum), device=device)
@@ -1600,13 +1596,11 @@
        # lazy setup_nccl_comm() call.
        with distributed_group(subgroup, trt_model) as migrated_model:
            with torch.no_grad():
                out_sub = migrated_model(inp)

-        _check_close(
-            out_sub, expected, f"[{label}] migrated to subgroup rank={rank}"
-        )
+        _check_close(out_sub, expected, f"[{label}] migrated to subgroup rank={rank}")

        # ---- Step 3: set_distributed_group (persistent, outside context) ----
        subgroup2 = dist.new_group(ranks=list(range(world_size)))
        torch_tensorrt.distributed.set_distributed_group(trt_model, subgroup2)
        # _state.pg is NOT set here — Python runtime falls back to world group
--- /home/runner/work/TensorRT/TensorRT/tools/llm/tensor_parallel_llama_multinode.py	2026-04-14 20:34:53.940401+00:00
+++ /home/runner/work/TensorRT/TensorRT/tools/llm/tensor_parallel_llama_multinode.py	2026-04-14 20:35:16.561193+00:00
@@ -194,11 +194,13 @@

            tokenizer = AutoTokenizer.from_pretrained(args.model)
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token

-            input_ids = tokenizer(args.prompt, return_tensors="pt")["input_ids"].to(DEVICE)
+            input_ids = tokenizer(args.prompt, return_tensors="pt")["input_ids"].to(
+                DEVICE
+            )
            max_len = input_ids.shape[1] + args.num_tokens

            logger.info("Running uncompiled PyTorch baseline ...")
            torch_tokens = generate(
                model, input_ids.clone(), max_len, tokenizer.eos_token_id
--- /home/runner/work/TensorRT/TensorRT/tools/llm/tensor_parallel_mixtral_llm.py	2026-04-14 20:34:53.940401+00:00
+++ /home/runner/work/TensorRT/TensorRT/tools/llm/tensor_parallel_mixtral_llm.py	2026-04-14 20:35:16.626752+00:00
@@ -135,16 +135,16 @@
            )
        cfg = model.config
        parallelize_module(model, device_mesh, build_tp_plan(cfg))

    cfg = model.config
-    assert cfg.num_key_value_heads % world_size == 0, (
-        f"num_key_value_heads ({cfg.num_key_value_heads}) not divisible by world_size ({world_size})"
-    )
-    assert cfg.num_attention_heads % world_size == 0, (
-        f"num_attention_heads ({cfg.num_attention_heads}) not divisible by world_size ({world_size})"
-    )
+    assert (
+        cfg.num_key_value_heads % world_size == 0
+    ), f"num_key_value_heads ({cfg.num_key_value_heads}) not divisible by world_size ({world_size})"
+    assert (
+        cfg.num_attention_heads % world_size == 0
+    ), f"num_attention_heads ({cfg.num_attention_heads}) not divisible by world_size ({world_size})"

    # After column-sharding Q/K/V, each rank holds num_heads // world_size
    # heads. Patch these so HuggingFace attention reshapes correctly.
    for layer in model.model.layers:
        layer.self_attn.num_heads = cfg.num_attention_heads // world_size
@@ -209,11 +209,11 @@
    parser.add_argument(
        "--sharded_checkpoint",
        type=str,
        default="",
        help="Path to DCP sharded checkpoint (e.g. /mnt/cluster-shared/mixtral_sharded). "
-             "If set, skips HF weight download and loads only this rank's shard.",
+        "If set, skips HF weight download and loads only this rank's shard.",
    )
    args = parser.parse_args()

    device_mesh = init_device_mesh("cuda", (world_size,))

--- /home/runner/work/TensorRT/TensorRT/tools/llm/utils.py	2026-04-14 20:34:53.941720+00:00
+++ /home/runner/work/TensorRT/TensorRT/tools/llm/utils.py	2026-04-14 20:35:17.028688+00:00
@@ -163,11 +163,13 @@
            # block-size padding produces s1*(8+s1-s1%8)>1 guards that the
            # symbolic solver can't verify without concrete values).  Without
            # bounds, dynamo traces symbolically and TRT infers the profile
            # from the first concrete shape it sees.
            torch._dynamo.mark_dynamic(input_seq, 1)
-        position_ids = torch.arange(input_seq.shape[1], device=input_seq.device).unsqueeze(0)
+        position_ids = torch.arange(
+            input_seq.shape[1], device=input_seq.device
+        ).unsqueeze(0)
        if dynamic_seqlen_range is not None:
            torch._dynamo.mark_dynamic(position_ids, 1)
        outputs = model(input_seq, position_ids=position_ids)
        logits = outputs.logits
        next_token_logits = logits[:, -1, :]

[github-actions]

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

--- /home/runner/work/TensorRT/TensorRT/tools/llm/tensor_parallel_mixtral_llm.py	2026-04-14 21:04:44.996426+00:00
+++ /home/runner/work/TensorRT/TensorRT/tools/llm/tensor_parallel_mixtral_llm.py	2026-04-14 21:05:11.648796+00:00
@@ -135,16 +135,16 @@
            )
        cfg = model.config
        parallelize_module(model, device_mesh, build_tp_plan(cfg))

    cfg = model.config
-    assert cfg.num_key_value_heads % world_size == 0, (
-        f"num_key_value_heads ({cfg.num_key_value_heads}) not divisible by world_size ({world_size})"
-    )
-    assert cfg.num_attention_heads % world_size == 0, (
-        f"num_attention_heads ({cfg.num_attention_heads}) not divisible by world_size ({world_size})"
-    )
+    assert (
+        cfg.num_key_value_heads % world_size == 0
+    ), f"num_key_value_heads ({cfg.num_key_value_heads}) not divisible by world_size ({world_size})"
+    assert (
+        cfg.num_attention_heads % world_size == 0
+    ), f"num_attention_heads ({cfg.num_attention_heads}) not divisible by world_size ({world_size})"

    # After column-sharding Q/K/V, each rank holds num_heads // world_size
    # heads. Patch these so HuggingFace attention reshapes correctly.
    for layer in model.model.layers:
        layer.self_attn.num_heads = cfg.num_attention_heads // world_size
@@ -209,11 +209,11 @@
    parser.add_argument(
        "--sharded_checkpoint",
        type=str,
        default="",
        help="Path to DCP sharded checkpoint (e.g. /mnt/cluster-shared/mixtral_sharded). "
-             "If set, skips HF weight download and loads only this rank's shard.",
+        "If set, skips HF weight download and loads only this rank's shard.",
    )
    args = parser.parse_args()

    device_mesh = init_device_mesh("cuda", (world_size,))