Native DistCollective fails at build time when fused into a Myelin ForeignNode (large multi-device graph)

Bug Description

The TRT-11 native multi-device collectives added in #4321 (add_dist_collective in py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py) work for small / standalone graphs, but fail during the TensorRT engine build on a large graph when a DistCollective layer is fused — together with surrounding reshape/unsqueeze/pointwise ops — into a Myelin ForeignNode. Myelin can't map the collective kind (toMyelinCommKind) once the IDistCollectiveLayer is absorbed into the fused node:

ERROR: Error Code 9: Skipping tactic 0x0 due to exception CollectiveOperation not supported for collective In toMyelinCommKind at /_src/optimizer/myelin/myelinDistCollectiveLayer.cpp:49
ERROR: Error Code 10: Internal Error (Could not find any implementation for node {ForeignNode[[UNSQUEEZE]...[SHUFFLE]-[.../proj_out/_reshape_copy]]} In computeCosts ... optimizer.cpp:4284)

Standalone/unfused collectives build and run fine (a 2-GPU all_reduce+all_to_all compiles + runs, no TRT-LLM); the failure only appears when the collective is fused into a ForeignNode, which happens naturally in a real 48-layer transformer (~39k nodes) with Ulysses all_to_all + tensor-parallel all_reduce.

To Reproduce (2 ranks)

A module with Linear → reshape/unsqueeze around a dist.all_to_all_single, so the collective is surrounded by shuffle/pointwise ops that Myelin fuses into one ForeignNode:

import torch, torch.nn as nn, torch.distributed as dist
import torch_tensorrt

class M(nn.Module):
    def __init__(self):
        super().__init__(); self.pin = nn.Linear(64,64); self.pout = nn.Linear(64,64)
    def forward(self, x):                       # [B,S,64]
        x = self.pin(x); b,s,h = x.shape
        x = x.reshape(b*s, h).unsqueeze(0).reshape(b*s, h).contiguous()
        y = torch.empty_like(x); dist.all_to_all_single(y, x)   # -> tensorrt::fused_nccl_all_to_all
        return self.pout(y.reshape(b, s, h))

def main():
    dist.init_process_group("nccl"); r = dist.get_rank(); torch.cuda.set_device(r)
    m = M().cuda().eval(); x = torch.randn(2, 8, 64, device="cuda")
    ep = torch.export.export(m, (x,))
    trt = torch_tensorrt.dynamo.compile(ep, inputs=[x], use_distributed_mode_trace=True)  # native collectives
    trt(x)

main()   # torchrun --nproc_per_node=2 repro.py

What works vs fails

• WORKS: standalone / unfused IDistCollectiveLayer (2-GPU all_reduce + all_to_all, no TRT-LLM), parity vs eager.
• FAILS: large graph where the collective is fused into a Myelin ForeignNode (errors above).

What we tried (did not help)

• optimization_level=0 — does not prevent the Myelin fusion.
• require_full_compilation=False — no effect; this is a build-time Myelin failure, not a converter-coverage partition fallback.
• There is no exposed torch-tensorrt setting to mark an IDistCollectiveLayer non-fusable, and no public TensorRT BuilderFlag / PreviewFeature / env var to disable Myelin fusion or keep collective layers standalone.

Expected behavior

Either (a) Myelin should support the collective kind inside a fused ForeignNode (extend toMyelinCommKind), or (b) the builder should not fuse an IDistCollectiveLayer into a ForeignNode (keep it standalone) so large multi-device graphs build, or (c) torch-tensorrt should expose a way to force a fusion boundary around add_dist_collective layers.

Environment

• torch-tensorrt 2.14.0.dev0 built from main@d2d459a (includes #4321 "TRT 11 MD Ops")
• TensorRT 11.0.0.114 / 11.1.0.106 ; PyTorch 2.14 nightly (cu130) ; 8×A100 ; MULTIDEVICE_RUNTIME_10_16 preview enabled

Additional context

Flagged as a "Note for maintainers (separate issue)" in #4380 (which fixes the orthogonal subgroup/nbRanks divisibility problem — independent of this). Current workaround: add the fused collective ops (torch.ops.tensorrt.fused_nccl_all_to_all.default, ...fused_nccl_all_reduce.default) to torch_executed_ops so they run in PyTorch and create partition boundaries, keeping the DiT compute in TensorRT — at a graph-break/perf cost.