dynamo: keep native NCCL collectives out of Myelin ForeignNode fusion (cast-boundary)

[pkisfaludi-nv]

DRAFT — do not merge (author reviews). Depends on #4380.

Summary

On top of #4380 (which routes native NCCL collectives at the op's process subgroup), this adds the remaining changes that let the sharded LTX-2.3 DiT lower to a single TensorRT engine with native in-engine collectives (no TRT-LLM, no ONNX, no NCCL plugin). All changes are in py/torch_tensorrt/dynamo/conversion/impl/nccl_ops.py.

The three changes

1. Myelin fusion cast-boundary — _coll_boundary(). Myelin fuses a native IDistCollectiveLayer into a ForeignNode and then fails with CollectiveOperation ... not supported In toMyelinCommKind (see Native DistCollective fails at build time when fused into a Myelin ForeignNode (large multi-device graph) #4381). Inserting a genuine dtype reformat (fp32 by default, _COLL_BOUNDARY_DTYPE) on the collective's input/output creates a precision-change boundary Myelin will not fuse across (identity / same-dtype casts get folded away — only a real reformat works). This lets the full 48-block DiT compile to one engine with collectives intact.
2. num_ranks before get_output(). layer.num_ranks must be set before layer.get_output(0) so Myelin can infer the collective's output shape (e.g. all_gather dim0 = in_dim0 * num_ranks). Upstream sets it after, tripping AllGather requires world_size (>1) to be set for shape inference ... Call myelinGraphSetWorldSize().
3. AG_VIA_ALLREDUCE=1 (opt-in). Even with (2), TRT-11 Myelin can reject a fused native ALL_GATHER shape (needs the graph world_size, which torch-tensorrt never sets at build). The opt-in workaround expresses all_gather as scatter-this-rank's-slice-into-a-zero-padded-buffer then ALL_REDUCE(SUM) — numerically identical, and ALL_REDUCE's output shape == input shape so no build-time world_size is needed. Used for the video→audio KV gather in CP.

Also adds _mark_collective_boundary() (env TRT_COLL_BOUNDARY) as a debug-mark helper.

Relationship to #4380 / #4381

• Stacked on dynamo: target native NCCL collectives at the op's process (sub)group #4380 (subgroup routing). This PR preserves dynamo: target native NCCL collectives at the op's process (sub)group #4380's _collective_group_ranks(group_name, ...) group resolution and is purely additive; the diff currently also shows dynamo: target native NCCL collectives at the op's process (sub)group #4380's commit and will shrink to just these changes once dynamo: target native NCCL collectives at the op's process (sub)group #4380 lands.
• Native DistCollective fails at build time when fused into a Myelin ForeignNode (large multi-device graph) #4381 (Myelin can't lower collectives inside a ForeignNode) is the root-cause bug; the cast-boundary here is the user-space mitigation. The durable fix is Myelin/TensorRT-side (per-subgroup communicator routing + toMyelinCommKind), delivered in the internal TensorRT/Myelin MRs (not in this repo).

Validation

• CP8: full 22B DiT → 1 TRT engine, trt_engine_submodules=1, torch_coll_tokens=0 (zero PyTorch-side collectives), ENGINE_RAN_OK; parity vs eager CP8 cos 0.99602 (video) / 0.99965 (audio) — PASS. (8× A100-SXM4-80GB.)
• CP4×TP2: single-engine cosine-correct end-to-end when combined with the internal TensorRT/Myelin per-subgroup-comm routing.

Notes

• The cast-boundary (fp32 default) and AG_VIA_ALLREDUCE are env-gated and default-safe (world-group / no-op) so CP8 and existing single-group paths are unchanged.
• Applied uniformly to the native DistCollective converters (all_gather / reduce_scatter / all_reduce / all_to_all and the scatter/gather variants) for consistency; the num_ranks reorder is a safe correctness fix for all.

[meta-cla]

Hi @pkisfaludi-nv!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[pkisfaludi-nv]

We reproduced the follow-up setCommunicator() failure reported by @apbose on
Qwen3-30B-A3B and found a TensorRT core bug in addition to the original
ForeignNode build failure from #4381.

The engine does not derive mIsMdEngine from the presence of signal/wait
scaffolding. At deserialization it checks the serialized runners' getNbRanks().
For a Myelin runner, both the DistCollective and Attention translators assigned
one graph-wide mNbRanks field directly. Therefore a rank-1 attention op
translated after a rank-2 collective could overwrite the value and make a
compute-heavy collective engine deserialize as single-device even though its
comm_allreduce instruction remained present.

The core fix accumulates the maximum rank count across all MD-aware operations:

mNbRanks = std::max(mNbRanks, params.nbRanks);
mNbRanks = std::max(mNbRanks, contextParallelDegree);

Internal TensorRT MRs created with git trt mr:

• TensorRT rel-11.2 !49080
• TensorRT main cherry-pick !49081

The focused TensorRT regression constructs
ALL_REDUCE(nbRanks=2) -> ATTENTION(nbRanks=1). The old assignment fails
engine.isMdEngine(); the max-rank fix passes on A100.

Related CP4xTP2 subgroup-routing MRs are:

• Myelin per-subgroup NCCL communicators !19723
• TensorRT standalone DistCollective subgroup routing !49040
• TensorRT fused-Myelin subgroup routing !49062

PR #4388's precision boundary is still useful: it prevents the unsupported
Myelin absorption that causes the original #4381 build failure. The new core MR
fixes the separate runtime MD-classification failure instead of relying on
isolation as a workaround.

I also added a Torch-TensorRT two-rank regression test in commit ed8d41a6d.
It verifies compute + native all-reduce compile into exactly one TRT engine,
that no collective escapes into the outer PyTorch graph, and eager parity.
Actual 2xA100 pytest result: 1 passed in 19.32s.