From 586fbf9ba03b621851499bb6957e718e93307949 Mon Sep 17 00:00:00 2001
From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Date: Wed, 3 Jun 2026 10:28:42 -0700
Subject: [PATCH 1/5] [nvbugs/6248757][fix] keep shared-expert all-reduce on
 main stream in multi_stream_moe

The multi_stream_moe transform moved the entire shared-expert subgraph,
including its trailing row-parallel all-reduce, onto the auxiliary CUDA
stream. At TP=8 on SM100 the shared-expert symm-mem MULTIMEM all-reduce
then ran concurrently with the routed-expert all-reduce on the main
stream; under monolithic (decode) CUDA-graph replay the two collectives
interleaved across ranks and silently corrupted the output, causing the
flaky GSM8K accuracy failure. TP=4 was unaffected because MULTIMEM is
disabled for world_size 4 on SM100, and MMLU passed because prefill does
not use the monolithic CUDA graph.

Overlap only the shared-expert GEMMs on the aux stream and run the
trailing all-reduce on the main stream, so no two collectives execute on
different streams concurrently. Non-collective shared branches keep the
previous behavior.

Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
---
 .../transform/library/multi_stream_moe.py     | 102 ++++++++++++++----
 tests/integration/test_lists/waives.txt       |   3 +-
 2 files changed, 85 insertions(+), 20 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py
index dc71bf1ff7a6..33ba897f916e 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py
+++ b/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
 
diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index dfd056b822ab..df2ba651fd54 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -17,7 +17,8 @@ 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::TestNemotronNanoV3::test_accuracy[bf16-4-trtllm] SKIP (https://nvbugs/6185150)
+accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-4-trtllm] SKIP (https://nvbugs/6185150)
 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)

From 697127667ce5d291e1f04eadb81e682a7ac267cc Mon Sep 17 00:00:00 2001
From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Date: Sat, 6 Jun 2026 23:01:13 -0700
Subject: [PATCH 2/5] [nvbugs/6248757][test] verify shared-expert all-reduce
 stays on main stream

Add a regression suite to test_multi_stream_moe.py covering the
multi_stream_moe placement fix:

- test_shared_expert_all_reduce_stays_on_main_stream: a trailing
  shared-expert all-reduce is rewired after end_aux/wait_aux so the
  collective runs on the main stream and consumes the synced aux output,
  never landing inside the aux-stream region.
- test_mid_branch_all_reduce_skips_transform: a collective in the middle
  of the shared-expert branch cannot be split off, so the transform must
  skip the node rather than move a collective to the aux stream.

Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
---
 .../custom_ops/test_multi_stream_moe.py       | 188 ++++++++++++++++++
 1 file changed, 188 insertions(+)

diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py
index 41c0214ac977..44e63a75b071 100644
--- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py
+++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py
@@ -46,6 +46,11 @@
   3. Numerical correctness — output matches eager reference within tolerance.
   4. CUDA graph compatibility — capture + replay produces correct output.
   5. Multi-layer stacking — multiple MoE layers handled independently.
+
+Additionally, a regression suite for nvbugs/6248757 verifies that a trailing
+shared-expert all-reduce is kept on the MAIN stream (only the GEMMs overlap on
+the aux stream), and that a collective in the middle of the shared-expert
+branch makes the transform skip the node entirely.
 """
 
 import torch
@@ -696,3 +701,186 @@ def test_tuple_fork_pattern_and_correctness():
     assert len(begin_nodes) == 1, f"Expected exactly one begin_aux node, got {len(begin_nodes)}"
     assert begin_nodes[0].args[0].target is torch.ops.auto_deploy.mock_tuple_fork_moe_test.default
     _assert_numerical_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda"))
+
+
+# ===================================================================
+# Tests — shared-expert all-reduce placement (regression for nvbugs/6248757)
+# ===================================================================
+#
+# At TP>=6 on SM100 the shared-expert row-parallel all-reduce was moved onto
+# the aux stream together with the rest of the shared-expert subgraph.  Under
+# monolithic (decode) CUDA-graph replay the shared-expert collective on the aux
+# stream then ran concurrently with the routed-expert collective on the main
+# stream; the two symm-mem MULTIMEM all-reduces interleaved across ranks and
+# silently corrupted the output.  The transform must therefore keep any
+# trailing shared-expert collective on the MAIN stream and overlap only the
+# shared-expert GEMMs on the aux stream.
+
+# Custom op packet used directly as the FX node target after export.
+_ALL_REDUCE_TARGET = torch.ops.auto_deploy.torch_dist_all_reduce.default
+
+
+def _node_order(gm):
+    """Map each node to its position in graph (execution) order."""
+    return {n: i for i, n in enumerate(gm.graph.nodes)}
+
+
+def _first_node_with_target(gm, target):
+    """Return the first ``call_function`` node whose target is *target* (or None)."""
+    for n in gm.graph.nodes:
+        if n.op == "call_function" and n.target is target:
+            return n
+    return None
+
+
+class _GatedMLPWithAllReduce(nn.Module):
+    """Shared expert whose row-parallel ``down_proj`` is followed by an all-reduce.
+
+    Mirrors the real DeepSeek/GLM4 sharded graph: the trailing collective (the
+    tensor-parallel shard reduction) is the last op of the shared-expert branch
+    and is what feeds the merge ``add``.
+    """
+
+    def __init__(self, hidden_dim: int, intermediate_dim: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
+        self.up_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
+        self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
+        return torch.ops.auto_deploy.torch_dist_all_reduce(y, "AUTO")
+
+
+class _MidBranchAllReduceMLP(nn.Module):
+    """Shared expert with an all-reduce in the *middle* of the branch.
+
+    The collective is not the trailing shared-output op, so it cannot be split
+    off onto the main stream — the transform must skip this node entirely
+    rather than place a collective on the aux stream.
+    """
+
+    def __init__(self, hidden_dim: int, intermediate_dim: int):
+        super().__init__()
+        self.up_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
+        self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = torch.relu(self.up_proj(x))
+        h = torch.ops.auto_deploy.torch_dist_all_reduce(h, "AUTO")
+        return self.down_proj(h)
+
+
+class MockSharedExpertAllReduceMoELayer(nn.Module):
+    """DeepSeek/GLM4 MoE layer whose shared expert ends in a row-parallel all-reduce."""
+
+    def __init__(self, hidden_dim: int, intermediate_dim: int, num_experts: int = 8):
+        super().__init__()
+        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
+        self.shared_experts = _GatedMLPWithAllReduce(hidden_dim, intermediate_dim)
+        self.expert_weight = nn.Parameter(torch.randn(hidden_dim, hidden_dim))
+        self.layernorm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        identity = hidden_states
+        logits = self.gate(hidden_states)
+        routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1)
+
+        moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test(
+            hidden_states, selected_experts, routing_weights, self.expert_weight
+        )
+        shared_out = self.shared_experts(identity)
+
+        return self.layernorm(moe_out + shared_out)
+
+
+class MockMidBranchAllReduceMoELayer(nn.Module):
+    """MoE layer whose shared expert has a collective in the middle of the branch."""
+
+    def __init__(self, hidden_dim: int, intermediate_dim: int, num_experts: int = 8):
+        super().__init__()
+        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
+        self.shared_experts = _MidBranchAllReduceMLP(hidden_dim, intermediate_dim)
+        self.expert_weight = nn.Parameter(torch.randn(hidden_dim, hidden_dim))
+        self.layernorm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        identity = hidden_states
+        logits = self.gate(hidden_states)
+        routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1)
+
+        moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test(
+            hidden_states, selected_experts, routing_weights, self.expert_weight
+        )
+        shared_out = self.shared_experts(identity)
+
+        return self.layernorm(moe_out + shared_out)
+
+
+def test_shared_expert_all_reduce_stays_on_main_stream():
+    """Regression for nvbugs/6248757.
+
+    A trailing shared-expert all-reduce must run on the MAIN stream: the
+    transform overlaps only the shared-expert GEMMs on the aux stream and emits
+    ``end_aux`` / ``wait_aux`` *before* the collective, so the collective is
+    rewired to consume the synced aux output and never lands inside the aux
+    region.
+    """
+    hidden_dim, intermediate_dim = 128, 256
+    cuda_stream_manager.add_device(torch.cuda.current_device())
+
+    model = MockSharedExpertAllReduceMoELayer(hidden_dim, intermediate_dim).eval().to("cuda")
+    example = torch.randn(4, hidden_dim, device="cuda")
+    gm = _build_gm(model, example)
+
+    gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS)
+
+    assert num == 1, f"Expected 1 replacement, got {num}"
+    _assert_stream_nodes_present(gm)
+
+    order = _node_order(gm)
+    ar_node = _first_node_with_target(gm, _ALL_REDUCE_TARGET)
+    begin_aux = _first_node_with_target(gm, begin_aux_stream_passthrough)
+    end_aux = _first_node_with_target(gm, end_aux_stream_passthrough)
+    wait_aux = _first_node_with_target(gm, wait_aux_stream_passthrough)
+
+    assert ar_node is not None, "shared-expert all-reduce node missing after transform"
+
+    # The collective must run after the aux region is closed (i.e. after the
+    # current stream is switched back to main) — never between begin/end_aux.
+    assert order[ar_node] > order[end_aux], "all-reduce must run after end_aux (main stream)"
+    assert order[ar_node] > order[wait_aux], "all-reduce must run after wait_aux (main stream)"
+    assert not (order[begin_aux] < order[ar_node] < order[end_aux]), (
+        "all-reduce was placed inside the aux-stream region"
+    )
+
+    # The collective is rewired to consume the synced aux output, proving the
+    # GEMM result produced on the aux stream is awaited before the main-stream
+    # collective reads it.
+    assert wait_aux in ar_node.all_input_nodes, (
+        "all-reduce must consume wait_aux output (synced aux result)"
+    )
+
+
+def test_mid_branch_all_reduce_skips_transform():
+    """A non-trailing shared-expert collective must make the transform skip the node.
+
+    Such a collective cannot be split off onto the main stream, so the transform
+    must bail out rather than move a collective to the aux stream.
+    """
+    hidden_dim, intermediate_dim = 128, 256
+    cuda_stream_manager.add_device(torch.cuda.current_device())
+
+    model = MockMidBranchAllReduceMoELayer(hidden_dim, intermediate_dim).eval().to("cuda")
+    example = torch.randn(4, hidden_dim, device="cuda")
+    gm = _build_gm(model, example)
+
+    gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS)
+
+    assert num == 0, f"Expected 0 replacements (mid-branch collective), got {num}"
+
+    # No aux-stream nodes should be inserted when the transform bails out.
+    targets = _stream_targets(gm)
+    assert begin_aux_stream_passthrough not in targets
+    assert end_aux_stream_passthrough not in targets
+    assert wait_aux_stream_passthrough not in targets

From 7bd171936c44a2019f829d9d719a82c663c1b5b6 Mon Sep 17 00:00:00 2001
From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Date: Sat, 6 Jun 2026 23:24:36 -0700
Subject: [PATCH 3/5] [nvbugs/6248757][test] add reviewer-requested
 trailing-all-reduce tests

Address review feedback (@tcherckez-nvidia) by adding the regression tests
from the reviewer's reference patch:

- New multi-GPU suite
  tests/.../multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py
  covering the real distributed path: a structural MPI check, an NCCL
  correctness check under CUDA-graph replay, and a SYMM_MEM corruption
  demo that builds the pre-fix (buggy) graph with the collective on the
  aux stream and shows it diverges from the fixed graph (skipped when
  SYMM_MEM/MULTIMEM is unavailable).

- Nemotron Ultra V3 single-GPU suite appended to
  tests/.../singlegpu/custom_ops/test_multi_stream_moe.py covering the
  fuse_rmsnorm_quant_nvfp4 topology (shared expert ending in a trailing
  all-reduce): graph-structure invariant, numerical correctness,
  CUDA-graph replay, and multi-layer stacking.

These complement the existing single-GPU placement tests and exercise the
fix on the multi-rank collective path it actually targets.

Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
---
 tests/integration/test_lists/waives.txt       |   2 -
 ...est_multi_stream_moe_trailing_allreduce.py | 541 ++++++++++++++++++
 .../custom_ops/test_multi_stream_moe.py       | 217 ++++++-
 3 files changed, 757 insertions(+), 3 deletions(-)
 create mode 100644 tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py

diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index df2ba651fd54..5b397d24f6f6 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -17,8 +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::TestNemotronNanoV3::test_accuracy[bf16-4-trtllm] SKIP (https://nvbugs/6185150)
-accuracy/test_llm_api_autodeploy.py::TestNemotronNanoV3::test_accuracy[nvfp4-4-trtllm] SKIP (https://nvbugs/6185150)
 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)
diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py
new file mode 100644
index 000000000000..ea6c223fe7e4
--- /dev/null
+++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py
@@ -0,0 +1,541 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Multi-GPU regression tests for nvbugs/6248757.
+
+Background
+----------
+``fuse_rmsnorm_quant_nvfp4`` restructures the shared-expert branch of a
+TP-sharded MoE layer so that a standalone ``trtllm_dist_all_reduce`` directly
+feeds the merge node::
+
+    hidden ─┬─ gate ─ topk ─── moe_fused ─── moe_out ───────────────────┐
+             └─ up_proj ─ relu² ─ down_proj ─ all_reduce ─ shared_out ─┴─ add
+
+Before PR #14917, ``_execute_shared_expert_in_aux_stream`` placed ``end_aux``
+*after* the all_reduce, putting the collective on the aux stream.  With
+AllReduceStrategy.SYMM_MEM (used in Nemotron Ultra V3 production config) two
+concurrent SYMM_MEM ops on different streams interleave across ranks under
+monolithic CUDA-graph replay and silently corrupt the output.
+
+Three test scenarios
+--------------------
+1. ``test_structural_multigpu`` — graph-level check: after the transform the
+   all_reduce must appear *after* ``end_aux``.  FAILS pre-PR, PASSES with fix.
+
+2. ``test_correctness_nccl_cuda_graph`` — NCCL correctness under CUDA graph.
+   NCCL serialises submissions CPU-side so the race does not manifest; serves
+   as a regression guard.
+
+3. ``test_corruption_symm_mem_cuda_graph`` — explicitly builds the *buggy*
+   graph, forces asymmetric submission ordering across ranks, and verifies the
+   buggy graph produces wrong output while the fixed graph is correct.
+   Skipped when SYMM_MEM is unavailable.
+"""
+
+import traceback
+
+import pytest
+import torch
+from torch.distributed import DistNetworkError
+
+# MPI pool leaks a thread on shutdown — suppress the threadleak warning.
+pytestmark = pytest.mark.threadleak(enabled=False)
+
+
+# ---------------------------------------------------------------------------
+# Worker helpers (everything torch.ops-related is inside workers to avoid
+# cloudpickle issues when serialising across MPI)
+# ---------------------------------------------------------------------------
+
+
+def _init_dist(port):
+    import torch.distributed as dist
+
+    import tensorrt_llm
+    import tensorrt_llm._torch.auto_deploy.custom_ops  # noqa: F401 — registers custom ops
+    from tensorrt_llm._torch.auto_deploy.distributed.common import initialize_or_skip
+    from tensorrt_llm._utils import get_free_port, mpi_broadcast
+
+    rank = tensorrt_llm.mpi_rank()
+    torch.cuda.set_device(rank)
+    # Rank 0 picks a free port and broadcasts it so all workers use the same one.
+    if port is None:
+        port = mpi_broadcast(get_free_port() if rank == 0 else None)
+    initialize_or_skip(port=port)
+    return rank, dist.get_world_size()
+
+
+def _cleanup():
+    import torch.distributed as dist
+
+    from tensorrt_llm._torch.auto_deploy.distributed.common import cleanup
+
+    if dist.is_initialized() and dist.get_world_size() > 1:
+        dist.barrier()
+    cleanup()
+
+
+def _make_model_and_example(hidden_dim, inter_dim, strategy, device="cuda"):
+    """Build the mock MoE layer with trailing all_reduce and an example input.
+
+    Defined inside worker functions to avoid cloudpickle capturing torch.ops
+    at module level.
+    """
+    import torch.nn as nn
+
+    # Register mock MoE op if not already registered.
+    op_name = "auto_deploy::mock_moe_trailing_ar"
+    if not hasattr(torch.ops.auto_deploy, "mock_moe_trailing_ar"):
+
+        @torch.library.custom_op(op_name, mutates_args=())
+        def _mock_moe(
+            x: torch.Tensor, sel: torch.Tensor, w: torch.Tensor, ew: torch.Tensor
+        ) -> torch.Tensor:
+            return torch.ops.aten.linear(x, ew)
+
+        @_mock_moe.register_fake
+        def _mock_moe_fake(x, sel, w, ew):
+            return torch.ops.aten.linear(x, ew)
+
+    moe_op = torch.ops.auto_deploy.mock_moe_trailing_ar
+    ar_op = torch.ops.auto_deploy.trtllm_dist_all_reduce
+
+    class _Layer(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.strategy = strategy
+            self.gate = nn.Linear(hidden_dim, 8, bias=False)
+            self.up = nn.Linear(hidden_dim, inter_dim, bias=False)
+            self.down = nn.Linear(inter_dim, hidden_dim, bias=False)
+            self.expert_w = nn.Parameter(torch.randn(hidden_dim, hidden_dim))
+            self.ln = nn.LayerNorm(hidden_dim)
+
+        def forward(self, x):
+            logits = self.gate(x)
+            rw, sel = torch.topk(logits, k=2, dim=-1)
+            shared = self.down(torch.relu(self.up(x)) ** 2)
+            shared_out = ar_op(shared, self.strategy)
+            moe_out = moe_op(x, sel, rw, self.expert_w)
+            return self.ln(shared_out + moe_out)
+
+    model = _Layer().eval().to(device)
+    example = torch.randn(4, hidden_dim, device=device)
+    moe_ops = [moe_op]
+    return model, example, moe_ops
+
+
+def _build_gm(model, example):
+    return torch.export.export(model, (example,)).module()
+
+
+# ---------------------------------------------------------------------------
+# Worker 1 — structural check
+# ---------------------------------------------------------------------------
+
+
+def _worker_structural(world_size, port):
+    import tensorrt_llm._torch.auto_deploy.custom_ops  # noqa: F401
+    from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_moe import (
+        _execute_shared_expert_in_aux_stream,
+    )
+    from tensorrt_llm._torch.auto_deploy.utils.multi_stream_utils import (
+        cuda_stream_manager,
+        end_aux_stream_passthrough,
+        wait_aux_stream_passthrough,
+    )
+    from tensorrt_llm._torch.auto_deploy.utils.node_utils import all_reduce_ops, is_op
+
+    rank, _ = _init_dist(port)
+    try:
+        cuda_stream_manager.add_device(rank)
+        model, example, moe_ops = _make_model_and_example(128, 256, "NCCL")
+        gm = _build_gm(model, example)
+        gm, num = _execute_shared_expert_in_aux_stream(gm, moe_ops)
+
+        assert num == 1, f"[rank {rank}] Expected 1 replacement, got {num}"
+
+        node_order = {n: i for i, n in enumerate(gm.graph.nodes)}
+        ar_ops = all_reduce_ops()
+        ar_node = next((n for n in gm.graph.nodes if is_op(n, ar_ops)), None)
+        end_aux_nodes = [
+            n
+            for n in gm.graph.nodes
+            if n.op == "call_function" and n.target is end_aux_stream_passthrough
+        ]
+        wait_aux_nodes = [
+            n
+            for n in gm.graph.nodes
+            if n.op == "call_function" and n.target is wait_aux_stream_passthrough
+        ]
+
+        assert ar_node is not None, f"[rank {rank}] No all_reduce node"
+        assert end_aux_nodes, f"[rank {rank}] No end_aux node"
+        assert wait_aux_nodes, f"[rank {rank}] No wait_aux node"
+
+        end_aux = end_aux_nodes[0]
+        wait_aux = wait_aux_nodes[0]
+
+        # Core invariant: collective must come AFTER the stream switch back to main.
+        assert node_order[ar_node] > node_order[end_aux], (
+            f"[rank {rank}] BUG: all_reduce before end_aux — on aux stream"
+        )
+        assert node_order[ar_node] > node_order[wait_aux], (
+            f"[rank {rank}] BUG: all_reduce before wait_aux"
+        )
+        assert end_aux.args[0] is not ar_node, (
+            f"[rank {rank}] BUG: end_aux wraps the all_reduce directly"
+        )
+        assert wait_aux in ar_node.all_input_nodes, (
+            f"[rank {rank}] wait_aux must feed the all_reduce"
+        )
+        return True
+    except Exception:
+        traceback.print_exc()
+        raise
+    finally:
+        _cleanup()
+
+
+# ---------------------------------------------------------------------------
+# Worker 2 — NCCL correctness under CUDA graph
+# ---------------------------------------------------------------------------
+
+
+def _worker_nccl_cuda_graph(world_size, port):
+    import tensorrt_llm._torch.auto_deploy.custom_ops  # noqa: F401
+    from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_moe import (
+        _execute_shared_expert_in_aux_stream,
+    )
+    from tensorrt_llm._torch.auto_deploy.utils.multi_stream_utils import cuda_stream_manager
+
+    rank, _ = _init_dist(port)
+    try:
+        cuda_stream_manager.add_device(rank)
+        torch.manual_seed(42 + rank)
+        model, example, moe_ops = _make_model_and_example(128, 256, "NCCL")
+        gm = _build_gm(model, example)
+        gm, num = _execute_shared_expert_in_aux_stream(gm, moe_ops)
+        assert num == 1
+
+        test_x = torch.randn(4, 128, device="cuda")
+        ref = model(test_x)
+
+        static_x = torch.randn_like(test_x)
+        static_out = torch.empty_like(ref)
+        for _ in range(3):
+            static_out.copy_(gm(static_x))
+
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            static_out.copy_(gm(static_x))
+
+        static_x.copy_(test_x)
+        g.replay()
+
+        assert torch.allclose(static_out, ref, atol=1e-4), (
+            f"[rank {rank}] CUDA graph mismatch: max diff {(static_out - ref).abs().max().item()}"
+        )
+        return True
+    except Exception:
+        traceback.print_exc()
+        raise
+    finally:
+        _cleanup()
+
+
+# ---------------------------------------------------------------------------
+# Worker 3 — SYMM_MEM corruption demo
+# ---------------------------------------------------------------------------
+
+
+def _worker_symm_mem_corruption(world_size, port):
+    """Build buggy and fixed graphs, run under SYMM_MEM + CUDA graph, compare."""
+    import torch.distributed as dist
+
+    import tensorrt_llm._torch.auto_deploy.custom_ops  # noqa: F401
+    from tensorrt_llm._torch.auto_deploy.custom_ops.distributed.trtllm_dist import (
+        is_trtllm_op_available,
+    )
+    from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_moe import (
+        _execute_shared_expert_in_aux_stream,
+        _find_merge_node,
+        _get_ancestors,
+    )
+    from tensorrt_llm._torch.auto_deploy.utils.multi_stream_utils import (
+        begin_aux_stream_passthrough,
+        cuda_stream_manager,
+        end_aux_stream_passthrough,
+        wait_aux_stream_passthrough,
+    )
+    from tensorrt_llm._torch.auto_deploy.utils.node_utils import all_reduce_ops, is_op
+    from tensorrt_llm.functional import AllReduce, AllReduceStrategy
+    from tensorrt_llm.mapping import Mapping
+
+    if not is_trtllm_op_available():
+        return "skip:no_trtllm_ops"
+
+    rank, wsize = _init_dist(port)
+    try:
+        cuda_stream_manager.add_device(rank)
+
+        # Check SYMM_MEM is available on this hardware.
+        try:
+            mapping = Mapping(world_size=wsize, tp_size=wsize, rank=rank)
+            ar_runner = AllReduce(
+                mapping=mapping, strategy=AllReduceStrategy.SYMM_MEM, dtype=torch.float16
+            )
+            if ar_runner.strategy != AllReduceStrategy.SYMM_MEM:
+                return "skip:no_symm_mem"
+        except Exception:
+            return "skip:no_symm_mem"
+
+        torch.manual_seed(42)
+        strategy = "SYMM_MEM"
+        hidden_dim, inter_dim = 128, 256
+        model, example, moe_ops = _make_model_and_example(hidden_dim, inter_dim, strategy)
+
+        ar_op = torch.ops.auto_deploy.trtllm_dist_all_reduce
+        ar_ops = all_reduce_ops()
+
+        # ----------------------------------------------------------------
+        # Build BUGGY graph: all_reduce placed on aux stream (pre-PR).
+        # ----------------------------------------------------------------
+        def make_buggy_gm():
+            gm = _build_gm(model, example)
+            graph = gm.graph
+            node_order_snap = {n: i for i, n in enumerate(graph.nodes)}
+
+            moe_node = next(n for n in graph.nodes if is_op(n, moe_ops))
+            merge_node = _find_merge_node(moe_node)
+            assert merge_node is not None
+
+            moe_anc = _get_ancestors(moe_node)
+            moe_anc.add(moe_node)
+
+            shared_output = routed_output = None
+            for arg in merge_node.all_input_nodes:
+                arg_anc = _get_ancestors(arg)
+                if moe_node in arg_anc or arg is moe_node:
+                    routed_output = arg
+                elif arg in moe_anc or arg.op != "call_function":
+                    pass
+                else:
+                    shared_output = arg
+
+            assert shared_output is not None and is_op(shared_output, ar_ops)
+
+            shared_nodes, fork_point, visited = [], None, set()
+            queue = [shared_output]
+            while queue:
+                n = queue.pop(0)
+                if n in visited:
+                    continue
+                visited.add(n)
+                if n.op == "get_attr":
+                    continue
+                if n in moe_anc:
+                    if fork_point is None or node_order_snap.get(n, 0) > node_order_snap.get(
+                        fork_point, 0
+                    ):
+                        fork_point = n
+                    continue
+                shared_nodes.append(n)
+                for inp in n.all_input_nodes:
+                    queue.append(inp)
+
+            shared_nodes.sort(key=lambda n: node_order_snap.get(n, 0))
+            first_shared = shared_nodes[0]
+
+            with graph.inserting_before(first_shared):
+                beg = graph.call_function(begin_aux_stream_passthrough, args=(fork_point,))
+            first_shared.args = tuple(beg if a is fork_point else a for a in first_shared.args)
+
+            # BUG: end_aux inserted AFTER the all_reduce → collective on aux stream.
+            with graph.inserting_after(shared_output):
+                end = graph.call_function(end_aux_stream_passthrough, args=(shared_output,))
+            merge_node.args = tuple(end if a is shared_output else a for a in merge_node.args)
+
+            with graph.inserting_before(merge_node):
+                wait = graph.call_function(wait_aux_stream_passthrough, args=(routed_output,))
+            merge_node.args = tuple(wait if a is routed_output else a for a in merge_node.args)
+
+            # Add second all_reduce on main stream (stands in for routed-expert AR).
+            out_node = next(n for n in reversed(list(graph.nodes)) if n.op == "output")
+            out_arg = out_node.args[0]
+            with graph.inserting_before(out_node):
+                second_ar = graph.call_function(ar_op.default, args=(out_arg, strategy))
+            out_node.args = (second_ar,)
+            graph.lint()
+            gm.recompile()
+            return gm
+
+        # ----------------------------------------------------------------
+        # Build FIXED graph: all_reduce on main stream.
+        # ----------------------------------------------------------------
+        def make_fixed_gm():
+            gm = _build_gm(model, example)
+            gm, num = _execute_shared_expert_in_aux_stream(gm, moe_ops)
+            assert num == 1
+            graph = gm.graph
+            out_node = next(n for n in reversed(list(graph.nodes)) if n.op == "output")
+            out_arg = out_node.args[0]
+            with graph.inserting_before(out_node):
+                second_ar = graph.call_function(ar_op.default, args=(out_arg, strategy))
+            out_node.args = (second_ar,)
+            graph.lint()
+            gm.recompile()
+            return gm
+
+        buggy_gm = make_buggy_gm()
+        fixed_gm = make_fixed_gm()
+
+        # ----------------------------------------------------------------
+        # Run both under CUDA graph with forced asymmetric stream ordering:
+        #   rank 0 → main stream waits for aux before second AR → aux submits first
+        #   rank 1 → no extra wait → main submits first
+        # SYMM_MEM collectives don't go through NCCL's CPU serialisation, so
+        # the cross-rank submission order mismatch produces wrong all_reduce
+        # results for the buggy graph.
+        # ----------------------------------------------------------------
+        aux_stream = cuda_stream_manager.get_stream(rank, "aux")
+        main_stream = cuda_stream_manager.get_stream(rank, "main")
+
+        def capture_and_replay(gm, x):
+            static_x = x.clone()
+            static_out = torch.empty_like(gm(static_x))
+            for _ in range(3):
+                static_out.copy_(gm(static_x))
+            # Asymmetric delay: rank 0 delays main before the second collective.
+            if rank == 0:
+                main_stream.wait_stream(aux_stream)
+            g = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(g):
+                static_out.copy_(gm(static_x))
+            static_x.copy_(x)
+            g.replay()
+            torch.cuda.synchronize()
+            return static_out.clone()
+
+        test_x = torch.randn(4, hidden_dim, device="cuda")
+        dist.barrier()
+
+        buggy_out = capture_and_replay(buggy_gm, test_x)
+
+        # Eager reference (sequential, no multi-stream).
+        with torch.cuda.stream(main_stream):
+            ref_out = model(test_x)
+            ref_out = ar_op(ref_out, strategy)
+        torch.cuda.synchronize()
+
+        dist.barrier()
+        fixed_out = capture_and_replay(fixed_gm, test_x)
+
+        # Fixed graph must be correct.
+        fixed_correct = torch.allclose(fixed_out, ref_out, atol=1e-3)
+        assert fixed_correct, (
+            f"[rank {rank}] Fixed graph wrong under SYMM_MEM CUDA graph: "
+            f"max diff = {(fixed_out - ref_out).abs().max().item():.4f}"
+        )
+
+        # Buggy graph should produce wrong output when SYMM_MEM is truly
+        # concurrent (interleaved) on this hardware.
+        buggy_correct = torch.allclose(buggy_out, ref_out, atol=1e-3)
+        if buggy_correct:
+            # SYMM_MEM may have serialised (e.g., fallback, world_size too small).
+            return "skip:race_not_triggered"
+
+        return True
+    except Exception:
+        traceback.print_exc()
+        raise
+    finally:
+        _cleanup()
+
+
+# ---------------------------------------------------------------------------
+# Pytest entry points — use MpiPoolSession.submit_sync like
+# test_allreduce_residual_rmsnorm_fusion.py to avoid cloudpickle torch.ops issues.
+# ---------------------------------------------------------------------------
+
+
+def _run_with_retries(worker_fn, world_size, **kwargs):
+    from tensorrt_llm.llmapi.mpi_session import MpiPoolSession
+
+    max_retries = 5
+    last_exc = None
+    for _ in range(max_retries):
+        pool = MpiPoolSession(n_workers=world_size)
+        try:
+            return pool.submit_sync(worker_fn, port=None, world_size=world_size, **kwargs)
+        except DistNetworkError as e:
+            last_exc = e
+            if "EADDRINUSE" not in str(e) and "address already in use" not in str(e).lower():
+                raise
+        finally:
+            pool.shutdown()
+    raise RuntimeError(f"Dist init failed after {max_retries} retries") from last_exc
+
+
+def _check_results(results):
+    """Assert all worker results are True; return first non-True for skip detection."""
+    for r in results:
+        if isinstance(r, str) and r.startswith("skip:"):
+            return r
+        assert r is True, f"Unexpected worker result: {r}"
+    return True
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires ≥ 2 GPUs")
+def test_structural_multigpu():
+    """Graph-level invariant holds in real multi-GPU MPI context.
+
+    FAILS on pre-PR code (all_reduce appears before end_aux → on aux stream).
+    PASSES on PR #14917 fix.
+    Uses NCCL — works on any multi-GPU setup.
+    """
+    results = _run_with_retries(_worker_structural, world_size=2)
+    _check_results(results)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires ≥ 2 GPUs")
+def test_correctness_nccl_cuda_graph():
+    """Fixed graph is numerically correct under NCCL + CUDA graph replay.
+
+    NCCL serialises collective submissions CPU-side, so the race between two
+    concurrent collectives does not manifest.  This test always passes on both
+    buggy and fixed code — it is a correctness regression guard.
+    """
+    results = _run_with_retries(_worker_nccl_cuda_graph, world_size=2)
+    _check_results(results)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires ≥ 2 GPUs")
+def test_corruption_symm_mem_cuda_graph():
+    """Buggy graph produces wrong output; fixed graph is correct — with SYMM_MEM.
+
+    Explicitly builds the buggy graph (all_reduce on aux stream), forces
+    asymmetric submission ordering across ranks, and verifies the wrong output.
+    Then runs the fixed graph and verifies correctness.
+
+    Skipped when SYMM_MEM is unavailable or the race does not manifest on this
+    hardware (e.g., world_size below MULTIMEM threshold).
+    """
+    results = _run_with_retries(_worker_symm_mem_corruption, world_size=2)
+    outcome = _check_results(results)
+    if isinstance(outcome, str) and outcome.startswith("skip:"):
+        pytest.skip(
+            f"SYMM_MEM race not reproducible ({outcome.split(':', 1)[1]}); "
+            f"try world_size ≥ 6 on SM100 for reliable MULTIMEM activation"
+        )
diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py
index 44e63a75b071..d60f97698a95 100644
--- a/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py
+++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/test_multi_stream_moe.py
@@ -50,9 +50,14 @@
 Additionally, a regression suite for nvbugs/6248757 verifies that a trailing
 shared-expert all-reduce is kept on the MAIN stream (only the GEMMs overlap on
 the aux stream), and that a collective in the middle of the shared-expert
-branch makes the transform skip the node entirely.
+branch makes the transform skip the node entirely.  A Nemotron Ultra V3 suite
+(produced by ``fuse_rmsnorm_quant_nvfp4`` at TP>=2) additionally checks graph
+structure, numerical correctness, CUDA-graph replay, and multi-layer stacking
+for the trailing-all-reduce topology.
 """
 
+from unittest.mock import patch
+
 import torch
 import torch.nn as nn
 
@@ -65,6 +70,7 @@
     end_aux_stream_passthrough,
     wait_aux_stream_passthrough,
 )
+from tensorrt_llm._torch.auto_deploy.utils.node_utils import all_reduce_ops, is_op
 
 # ---------------------------------------------------------------------------
 # Mock fused-MoE custom op (distinct name to avoid conflicts with other tests)
@@ -884,3 +890,212 @@ def test_mid_branch_all_reduce_skips_transform():
     assert begin_aux_stream_passthrough not in targets
     assert end_aux_stream_passthrough not in targets
     assert wait_aux_stream_passthrough not in targets
+
+
+# ===================================================================
+# Tests — Nemotron Ultra V3: shared expert with trailing all-reduce
+# (nvbugs/6248757 — the bug introduced by fuse_rmsnorm_quant_nvfp4)
+# ===================================================================
+
+
+class MockNemotronUltraSharedAllReduceMoELayer(nn.Module):
+    """Shared expert branch ending in an all-reduce collective.
+
+    Simulates the FX graph topology produced by ``fuse_rmsnorm_quant_nvfp4``
+    on Nemotron Ultra V3 at TP>=2.  After that fusion the shared-expert
+    ``down_proj`` (row-wise TP split) emits a standalone
+    ``torch_dist_all_reduce`` whose output feeds the merge ``add`` directly::
+
+        hidden_states ─┬─ gate ─ topk ──── mock_fused_moe ──── moe_out ──────────┐
+                       └─ up_proj ─ relu² ─ down_proj ─ all_reduce ─ shared_out ─┴─ add ─ out
+
+    Before PR #14917 (buggy): ``_execute_shared_expert_in_aux_stream`` treats
+    ``shared_output`` (the all_reduce node) as the last shared op and inserts
+    ``end_aux`` *after* it.  The collective therefore executes on the aux stream.
+
+    After PR #14917 (fixed): the all_reduce is detected as a collective, split
+    off from the aux region, and forced to run on the main stream with
+    ``end_aux`` / ``wait_aux`` inserted before it.
+    """
+
+    def __init__(self, hidden_dim: int, intermediate_dim: int, num_experts: int = 8):
+        super().__init__()
+        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
+        self.up_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
+        self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False)
+        self.expert_weight = nn.Parameter(torch.randn(hidden_dim, hidden_dim))
+        self.layernorm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residuals = hidden_states
+        logits = self.gate(hidden_states)
+        routing_weights, selected_experts = torch.topk(logits, k=2, dim=-1)
+
+        # Shared expert path: GEMMs on aux stream, all_reduce on main stream.
+        shared = self.down_proj(torch.relu(self.up_proj(residuals)) ** 2)
+        shared_out = torch.ops.auto_deploy.torch_dist_all_reduce(shared, "sum")
+
+        moe_out = torch.ops.auto_deploy.mock_fused_moe_moe_test(
+            hidden_states, selected_experts, routing_weights, self.expert_weight
+        )
+        return self.layernorm(shared_out + moe_out)
+
+
+def _get_node_order(gm):
+    return {n: i for i, n in enumerate(gm.graph.nodes)}
+
+
+def _find_nodes_by_target(gm, target):
+    return [n for n in gm.graph.nodes if n.op == "call_function" and n.target is target]
+
+
+def _find_allreduce_node(gm):
+    ar_ops = all_reduce_ops()
+    return next((n for n in gm.graph.nodes if is_op(n, ar_ops)), None)
+
+
+def _assert_allreduce_on_main_stream(gm):
+    """The all_reduce must come AFTER end_aux and wait_aux in graph order.
+
+    This is the key structural invariant of the fix (PR #14917).
+
+    With the bug:  graph order is  begin_aux → ... → all_reduce → end_aux → merge
+                   (collective runs on aux stream)
+    With the fix:  graph order is  begin_aux → ... → end_aux → wait_aux → all_reduce → merge
+                   (collective runs on main stream)
+    """
+    node_order = _get_node_order(gm)
+    all_reduce_node = _find_allreduce_node(gm)
+    assert all_reduce_node is not None, "No all_reduce node found in graph"
+
+    end_aux_nodes = _find_nodes_by_target(gm, end_aux_stream_passthrough)
+    wait_aux_nodes = _find_nodes_by_target(gm, wait_aux_stream_passthrough)
+    assert end_aux_nodes, "end_aux_stream_passthrough not found"
+    assert wait_aux_nodes, "wait_aux_stream_passthrough not found"
+
+    end_aux_node = end_aux_nodes[0]
+    wait_aux_node = wait_aux_nodes[0]
+
+    # Core invariant: all_reduce must come AFTER end_aux (and wait_aux) in
+    # graph order.  Violation means the collective is on the aux stream.
+    assert node_order[all_reduce_node] > node_order[end_aux_node], (
+        "BUG: all_reduce collective appears before end_aux in graph order — "
+        "it is running on the aux stream.  The collective must be on the main stream."
+    )
+    assert node_order[all_reduce_node] > node_order[wait_aux_node], (
+        "BUG: all_reduce appears before wait_aux — main stream is not waiting "
+        "for aux compute before running the collective."
+    )
+
+    # end_aux must NOT wrap the all_reduce node directly.
+    assert end_aux_node.args[0] is not all_reduce_node, (
+        "BUG: end_aux_stream_passthrough wraps the all_reduce node as its input, "
+        "meaning the collective ran on the aux stream."
+    )
+
+    # wait_aux must feed the all_reduce directly (it replaced the aux_boundary arg).
+    assert wait_aux_node in all_reduce_node.all_input_nodes, (
+        "wait_aux_stream_passthrough must be a direct input to the all_reduce node "
+        "so the main stream synchronizes with the aux stream before the collective."
+    )
+
+
+def test_trailing_allreduce_graph_structure():
+    """The all_reduce must stay on the main stream after the transform.
+
+    This test captures the exact bug from nvbugs/6248757:
+    - FAILS on pre-PR code: end_aux is inserted after the all_reduce, placing
+      the collective on the aux stream.
+    - PASSES on PR #14917 fix: end_aux/wait_aux are inserted before the
+      all_reduce, which runs on the main stream.
+    """
+    hidden_dim, intermediate_dim = 128, 256
+    cuda_stream_manager.add_device(torch.cuda.current_device())
+
+    model = MockNemotronUltraSharedAllReduceMoELayer(hidden_dim, intermediate_dim).eval().to("cuda")
+    example = torch.randn(4, hidden_dim, device="cuda")
+    gm = _build_gm(model, example)
+
+    gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS)
+
+    assert num == 1, f"Expected 1 replacement, got {num}"
+    _assert_stream_nodes_present(gm)
+    _assert_allreduce_on_main_stream(gm)
+
+
+def test_trailing_allreduce_correctness():
+    """Numerical correctness for the shared-expert all_reduce topology.
+
+    Patches torch.distributed.all_reduce to identity (correct for world_size=1)
+    so the test runs without an initialized process group.
+    """
+    hidden_dim, intermediate_dim = 128, 256
+    cuda_stream_manager.add_device(torch.cuda.current_device())
+
+    model = MockNemotronUltraSharedAllReduceMoELayer(hidden_dim, intermediate_dim).eval().to("cuda")
+    example = torch.randn(4, hidden_dim, device="cuda")
+    gm = _build_gm(model, example)
+
+    gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS)
+    assert num == 1
+
+    test_x = torch.randn(4, hidden_dim, device="cuda")
+    with patch("torch.distributed.all_reduce", lambda t, **kw: None):
+        ref = model(test_x)
+        out = gm(test_x)
+
+    assert torch.allclose(out, ref, atol=1e-5), (
+        f"Output mismatch after transform: max diff = {(out - ref).abs().max().item()}"
+    )
+
+
+def test_trailing_allreduce_cuda_graph():
+    """CUDA graph capture + replay for the shared-expert all_reduce topology."""
+    hidden_dim, intermediate_dim = 128, 256
+    cuda_stream_manager.add_device(torch.cuda.current_device())
+
+    model = MockNemotronUltraSharedAllReduceMoELayer(hidden_dim, intermediate_dim).eval().to("cuda")
+    example = torch.randn(4, hidden_dim, device="cuda")
+    gm = _build_gm(model, example)
+
+    gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS)
+    assert num == 1
+
+    with patch("torch.distributed.all_reduce", lambda t, **kw: None):
+        _assert_cuda_graph_correctness(gm, model, torch.randn(4, hidden_dim, device="cuda"))
+
+
+def test_trailing_allreduce_multi_layer():
+    """Two stacked layers with trailing all_reduce — both transformed, both correct."""
+    hidden_dim, intermediate_dim = 128, 256
+    cuda_stream_manager.add_device(torch.cuda.current_device())
+
+    model = (
+        nn.Sequential(
+            MockNemotronUltraSharedAllReduceMoELayer(hidden_dim, intermediate_dim),
+            MockNemotronUltraSharedAllReduceMoELayer(hidden_dim, intermediate_dim),
+        )
+        .eval()
+        .to("cuda")
+    )
+    example = torch.randn(4, hidden_dim, device="cuda")
+    gm = _build_gm(model, example)
+
+    gm, num = _execute_shared_expert_in_aux_stream(gm, _MOE_OPS)
+
+    assert num == 2, f"Expected 2 replacements, got {num}"
+    _assert_stream_nodes_present(gm)
+
+    # Both all_reduce nodes must be on the main stream.
+    node_order = _get_node_order(gm)
+    ar_ops = all_reduce_ops()
+    ar_nodes = [n for n in gm.graph.nodes if is_op(n, ar_ops)]
+    end_aux_nodes = _find_nodes_by_target(gm, end_aux_stream_passthrough)
+    assert len(ar_nodes) == 2, f"Expected 2 all_reduce nodes, got {len(ar_nodes)}"
+    assert len(end_aux_nodes) == 2, f"Expected 2 end_aux nodes, got {len(end_aux_nodes)}"
+    for ar_node in ar_nodes:
+        # Each all_reduce must come after at least one end_aux node.
+        assert any(node_order[ar_node] > node_order[e] for e in end_aux_nodes), (
+            f"all_reduce node {ar_node.name} is not after any end_aux — "
+            "it would run on the aux stream."
+        )

From 46754f6985525cf596540f729c364866958f3b77 Mon Sep 17 00:00:00 2001
From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Date: Mon, 8 Jun 2026 01:11:29 -0700
Subject: [PATCH 4/5] [nvbugs/6248757][fix] guard all_reduce/all_gather op
 lookup in standalone mode

The new shared-expert all-reduce gating in _execute_shared_expert_in_aux_stream
calls all_reduce_ops(), which referenced torch.ops.auto_deploy.trtllm_dist_all_reduce
unconditionally.  In the standalone llmc package the trtllm_dist module fails
to import (no tensorrt_llm available) and the op is never registered, so the
attribute access raises AttributeError and breaks every test that runs the
multi-stream-MoE transform (test_deepseek_glm4_*, test_nemotron_h_*,
test_fused_merge_*, test_tuple_fork_*).

Use the existing _auto_deploy_op helper to look up each op by name and silently
drop entries that did not register.  Behavior is unchanged when both backends
register; in standalone mode only the torch_dist variant is included, which
matches what is actually available.

Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
---
 .../_torch/auto_deploy/utils/node_utils.py    | 31 +++++++++++++------
 1 file changed, 22 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index f80d55f18870..17b4012a83e4 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/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
     )
 
 

From d3f00097f4bcba4e97f6dddb23b9bac5549bffbe Mon Sep 17 00:00:00 2001
From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
Date: Mon, 8 Jun 2026 05:54:34 -0700
Subject: [PATCH 5/5] [nvbugs/6248757][fix] import AllReduce from
 _torch.distributed

The trailing-all-reduce multi-GPU test imported AllReduce / AllReduceStrategy
from tensorrt_llm.functional, but functional.py only exposes the strategy
enums (AllReduceStrategy / AllReduceFusionOp / AllReduceParams). The runtime
class AllReduce lives in tensorrt_llm._torch.distributed; importing it from
the wrong module raises ImportError at collection time and skips every test
case in the file. Switch to the canonical import path used elsewhere in the
test tree (tests/unittest/_torch/multi_gpu/test_allreduce.py et al.).

Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com>
---
 .../custom_ops/test_multi_stream_moe_trailing_allreduce.py      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py b/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py
index ea6c223fe7e4..139d76cd5a97 100644
--- a/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py
+++ b/tests/unittest/auto_deploy/multigpu/custom_ops/test_multi_stream_moe_trailing_allreduce.py
@@ -280,7 +280,7 @@ def _worker_symm_mem_corruption(world_size, port):
         wait_aux_stream_passthrough,
     )
     from tensorrt_llm._torch.auto_deploy.utils.node_utils import all_reduce_ops, is_op
-    from tensorrt_llm.functional import AllReduce, AllReduceStrategy
+    from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy
     from tensorrt_llm.mapping import Mapping
 
     if not is_trtllm_op_available():