[https://nvbugs/6248757][fix] Avoid running all reduce in aux stream

tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py

@@ -29,7 +29,7 @@
     end_aux_stream_passthrough,
     wait_aux_stream_passthrough,
 )
-from ...utils.node_utils import has_shape, is_op
+from ...utils.node_utils import all_reduce_ops, has_shape, is_op
 from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry
 
 
@@ -194,7 +194,32 @@ def _execute_shared_expert_in_aux_stream(
 
         # Order shared nodes by their position in the graph.
         shared_nodes.sort(key=lambda n: node_order.get(n, 0))
-        first_shared = shared_nodes[0]
+
+        # Collectives (all-reduce) in the shared-expert branch must stay on the
+        # MAIN stream.  A collective synchronizes across ranks, and that
+        # rendezvous does not compose with per-rank aux-stream overlap: when the
+        # shared-expert all-reduce is captured on the aux stream while the
+        # routed-expert all-reduce runs on the main stream, the two symm-mem
+        # MULTIMEM collectives (world_size >= 6 on SM100) interleave across ranks
+        # under monolithic CUDA-graph replay and silently corrupt the output.
+        # We therefore overlap only the shared-expert GEMMs on the aux stream and
+        # run the trailing all-reduce on the main stream.
+        ar_ops = all_reduce_ops()
+        collective_node = shared_output if is_op(shared_output, ar_ops) else None
+        aux_region = [n for n in shared_nodes if n is not collective_node]
+
+        # The aux-stream region must contain compute and must not itself contain
+        # a collective (only a trailing shared-output collective can be split
+        # off safely).
+        if not aux_region or any(is_op(n, ar_ops) for n in aux_region):
+            ad_logger.warning(
+                f"Shared-expert branch of MoE node {moe_node.name} has no aux-stream "
+                "compute outside of a collective; skipping multi-stream transform for "
+                "this node."
+            )
+            continue
+
+        first_shared = aux_region[0]
 
         # Sanity check: the first shared op must directly consume the fork
         # point so we can wire begin_aux_stream_passthrough into it.
@@ -222,28 +247,67 @@ def _execute_shared_expert_in_aux_stream(
             begin_aux_node if arg is fork_point else arg for arg in first_shared.args
         )
 
-        # ---- Step 5: Insert end_aux after the last shared-expert op. ----
-        with graph.inserting_after(shared_output):
-            end_aux_node = graph.call_function(
-                end_aux_stream_passthrough,
-                args=(shared_output,),
+        if collective_node is None:
+            # ---- Step 5: Insert end_aux after the last shared-expert op. ----
+            with graph.inserting_after(shared_output):
+                end_aux_node = graph.call_function(
+                    end_aux_stream_passthrough,
+                    args=(shared_output,),
+                )
+
+            # Replace shared-expert input to the merge node with end_aux output.
+            merge_node.args = tuple(
+                end_aux_node if arg is shared_output else arg for arg in merge_node.args
             )
 
-        # Replace shared-expert input to the merge node with end_aux output.
-        merge_node.args = tuple(
-            end_aux_node if arg is shared_output else arg for arg in merge_node.args
-        )
+            # ---- Step 6: Insert wait_aux before the merge node. ----
+            with graph.inserting_before(merge_node):
+                wait_aux_node = graph.call_function(
+                    wait_aux_stream_passthrough,
+                    args=(routed_output,),
+                )
 
-        # ---- Step 6: Insert wait_aux before the merge node. ----
-        with graph.inserting_before(merge_node):
-            wait_aux_node = graph.call_function(
-                wait_aux_stream_passthrough,
-                args=(routed_output,),
+            merge_node.args = tuple(
+                wait_aux_node if arg is routed_output else arg for arg in merge_node.args
+            )
+        else:
+            # The trailing all-reduce stays on the main stream.  End the aux
+            # region after the last aux-stream compute op (e.g. the rowwise
+            # down-projection) and make the main stream wait for it before the
+            # collective consumes the result.
+            aux_boundary = max(
+                (a for a in collective_node.all_input_nodes if a in aux_region),
+                key=lambda n: node_order.get(n, 0),
+                default=None,
             )
+            if aux_boundary is None:
+                ad_logger.warning(
+                    f"Could not find aux-stream input to the shared-expert collective "
+                    f"for MoE node {moe_node.name}; skipping multi-stream transform."
+                )
+                continue
 
-        merge_node.args = tuple(
-            wait_aux_node if arg is routed_output else arg for arg in merge_node.args
-        )
+            # ---- Step 5: end_aux after the last aux-stream op, switching the
+            # current stream back to main before the collective. ----
+            with graph.inserting_after(aux_boundary):
+                end_aux_node = graph.call_function(
+                    end_aux_stream_passthrough,
+                    args=(aux_boundary,),
+                )
+
+            # ---- Step 6: wait_aux so the main stream waits for the aux compute
+            # before running the collective on the main stream. ----
+            with graph.inserting_after(end_aux_node):
+                wait_aux_node = graph.call_function(
+                    wait_aux_stream_passthrough,
+                    args=(end_aux_node,),
+                )
+
+            # The collective now consumes the synced aux output and runs on the
+            # main stream; the merge node continues to consume the collective.
+            collective_node.args = tuple(
+                wait_aux_node if arg is aux_boundary else arg for arg in collective_node.args
+            )
 
         num_replaced += 1
 

tensorrt_llm/_torch/auto_deploy/utils/node_utils.py

@@ -808,23 +808,36 @@ def all_gather_ops() -> frozenset:
 
     Strategy (AUTO/SYMM_MEM) and workspace_id (for symm-mem ProcessGroup
     selection) flow through as op arguments, not as separate op identities.
+
+    The TRT-LLM-backed ops are silently skipped if their custom_ops module
+    failed to register (e.g. in the standalone ``llmc`` package, where
+    ``trtllm_dist`` is not importable).
     """
     return frozenset(
-        {
-            torch.ops.auto_deploy.trtllm_dist_all_gather,
-            torch.ops.auto_deploy.torch_dist_all_gather,
-        }
+        op
+        for op in (
+            _auto_deploy_op("trtllm_dist_all_gather"),
+            _auto_deploy_op("torch_dist_all_gather"),
+        )
+        if op is not None
     )
 
 
 @functools.cache
 def all_reduce_ops() -> frozenset:
-    """All AllReduce custom op packets recognized by AutoDeploy."""
+    """All AllReduce custom op packets recognized by AutoDeploy.
+
+    The TRT-LLM-backed op is silently skipped if its custom_ops module
+    failed to register (e.g. in the standalone ``llmc`` package, where
+    ``trtllm_dist`` is not importable).
+    """
     return frozenset(
-        {
-            torch.ops.auto_deploy.trtllm_dist_all_reduce,
-            torch.ops.auto_deploy.torch_dist_all_reduce,
-        }
+        op
+        for op in (
+            _auto_deploy_op("trtllm_dist_all_reduce"),
+            _auto_deploy_op("torch_dist_all_reduce"),
+        )
+        if op is not None
     )
 
 

tests/integration/test_lists/waives.txt

@@ -17,7 +17,6 @@ accuracy/test_llm_api_autodeploy.py::TestGemmaE2B::test_gemma4_e2b_it SKIP (http
 accuracy/test_llm_api_autodeploy.py::TestMiniMaxM2::test_finegrained_fp8 SKIP (https://nvbugs/6158397)
 accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[deepseek-ai_DeepSeek-R1-0528-True] SKIP (https://nvbugs/6278380)
 accuracy/test_llm_api_autodeploy.py::TestModelRegistryAccuracy::test_autodeploy_from_registry[nvidia_Llama-3.1-8B-Instruct-NVFP4-True] SKIP (https://nvbugs/6245279)
-accuracy/test_llm_api_autodeploy.py::TestNemotronUltraV3::test_accuracy[nvfp4-8] SKIP (https://nvbugs/6248757)
 accuracy/test_llm_api_autodeploy.py::TestNemotronV2::test_fp8[True] SKIP (https://nvbugs/6261164)
 accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_bf16_small[4] SKIP (https://nvbugs/6158397)
 accuracy/test_llm_api_autodeploy.py::TestQwen3_5_397B_MoE::test_nvfp4[8] SKIP (https://nvbugs/6278380)