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

tensorrt-cicd · 2026-06-03T17:47:25Z

Summary

Root cause: At TP=8 on B200 (SM100), the shared-expert symm-mem MULTIMEM all-reduce executed on the aux stream concurrently with the routed-expert all-reduce on the main stream. Under the decode-time monolithic CUDA graph, the two cross-rank collectives interleave across ranks during replay and silently corrupt results, producing the flaky GSM8K failure.
Fix: Keep the optimization enabled. Detect a trailing all-reduce in the shared-expert branch (all_reduce_ops()); when present, overlap only the shared-expert GEMMs on the aux stream and emit end_aux/wait_aux before the collective so the all-reduce runs on the main stream. No two collectives ever run on different streams concurrently. If a collective appears mid-branch (cannot be split off), the node is skipped with a warning. Non-collective shared branches retain the original behavior. Verified at TP=8 with the feature ON (matches=48/rank): MMLU 86.891, GSM8K 94.141 (≥ 81.517), 1 passed. Waiver waives.txt line for nvbugs/6248757 removed.
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6248757

Summary by CodeRabbit

Bug Fixes
- Enhanced multi-stream Mixture of Experts (MoE) transform in auto-deployment to properly handle shared-expert branches with collective operations. This improvement ensures correct synchronization and routing efficiency when deploying complex MoE models.
Tests
- Re-enabled accuracy validation test for specific model configurations that was previously waived, ensuring comprehensive test coverage.

coderabbitai · 2026-06-03T17:51:45Z

📝 Walkthrough

Walkthrough

This PR extends the multi-stream MoE transform to correctly handle shared-expert branches with trailing all-reduce collectives. It detects collective operations, constrains aux-stream work to the non-collective portion, and inserts synchronization barriers between auxiliary and main streams. A previously-waived integration test is unwaived to validate the fix.

Changes

Collective-aware multi-stream MoE transform

Layer / File(s)	Summary
Collective detection and conditional synchronization `tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`	Import `all_reduce_ops` utility. Detect whether `shared_output` is an all-reduce collective and define aux-stream region as the non-collective portion. Rework end_aux/wait_aux wiring into conditional paths: without a collective, preserve prior behavior (end aux after `shared_output`, wait before merge); with a collective, find the latest aux-stream input feeding that collective, end aux after that input, insert `wait_aux_stream_passthrough`, and rewire the collective to consume the waited output so synchronization occurs before the collective executes on the main stream. Skip transform if aux region is empty or contains collectives.
Test waiver removal `tests/integration/test_lists/waives.txt`	Remove waived integration-test entry for `TestNemotronUltraV3::test_accuracy[nvfp4-8]`, allowing the test to run and validate that the collective-aware MoE transform fix handles this case correctly.

Sequence Diagram

sequenceDiagram
  participant MainStream
  participant AuxStream
  participant Collective as shared_output<br/>(all-reduce)
  participant Merge as merge_output
  
  AuxStream->>AuxStream: compute on aux-stream<br/>(non-collective nodes)
  AuxStream->>AuxStream: end_aux_stream_passthrough
  MainStream->>MainStream: route + compute
  MainStream->>Collective: wait_aux_stream_passthrough<br/>(sync before collective)
  Collective->>Collective: all-reduce collective<br/>(executes on main stream)
  Collective->>Merge: collective output
  MainStream->>Merge: routed output
  Merge->>Merge: merge result

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#14854: Both PRs modify tests/integration/test_lists/waives.txt by changing which integration tests are waived.
NVIDIA/TensorRT-LLM#14839: Both PRs modify tests/integration/test_lists/waives.txt by changing SKIP/waiver entries.

Suggested reviewers

tcherckez-nvidia
greg-kwasniewski1

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description check	✅ Passed	The PR description includes root cause, solution, verification results, and test plan, covering all essential sections of the template.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title accurately describes the main fix: preventing all-reduce operations from running in the auxiliary stream, which is the core bug being addressed.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

suyoggupta · 2026-06-05T17:14:02Z

can we add a multi-gpu unit test that checks all_reduce is always placed on the main stream @tcherckez-nvidia

tcherckez-nvidia · 2026-06-07T05:44:19Z

Add tests, use this file as reference
tests_nvbugs6248757.patch

tcherckez-nvidia · 2026-06-07T06:30:37Z

/bot run --stage-list "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1"

tensorrt-cicd · 2026-06-07T06:36:51Z

PR_Github #52556 [ run ] triggered by Bot. Commit: 1e5386b Link to invocation

tensorrt-cicd · 2026-06-07T08:52:42Z

PR_Github #52556 [ run ] completed with state SUCCESS. Commit: 1e5386b
/LLM/main/L0_MergeRequest_PR pipeline #41841 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

tcherckez-nvidia · 2026-06-07T08:58:30Z

/bot run

tensorrt-cicd · 2026-06-07T09:04:05Z

PR_Github #52565 [ run ] triggered by Bot. Commit: 1e5386b Link to invocation

tensorrt-cicd · 2026-06-07T12:07:18Z

PR_Github #52565 [ run ] completed with state SUCCESS. Commit: 1e5386b
/LLM/main/L0_MergeRequest_PR pipeline #41849 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

tcherckez-nvidia · 2026-06-07T12:15:42Z

/bot run --reuse-test

tcherckez-nvidia · 2026-06-07T15:09:54Z

added tests dont pass, see root cause:

Likely cause: The PR added new tests to test_multi_stream_moe.py (the test_trailing_allreduce_, test_shared_expert_, test_mid_branch_* group) using torch.ops.auto_deploy.torch_dist_all_reduce, while the pre-existing tests in that same file (test_deepseek_glm4_, test_nemotron_h_, test_fused_merge_, test_tuple_fork_) use the op name trtllm_dist_all_reduce. Simultaneously, the new multi-GPU test file test_multi_stream_moe_trailing_allreduce.py references torch.ops.auto_deploy.trtllm_dist_all_reduce in _make_model_and_example(). This is a naming inconsistency: either the op was recently renamed from trtllm_dist_all_reduce → torch_dist_all_reduce without updating the pre-existing tests, or the new tests used the wrong (new) name while the standalone package doesn't expose the old alias. The standalone runner uses a different environment where the old name is absent, exposing the mismatch.

tcherckez-nvidia · 2026-06-08T04:37:17Z

/bot run --reuse-test

tensorrt-cicd · 2026-06-08T04:42:41Z

PR_Github #52657 [ run ] triggered by Bot. Commit: a8ffede Link to invocation

tensorrt-cicd · 2026-06-08T07:31:03Z

PR_Github #52657 [ run ] completed with state FAILURE. Commit: a8ffede
/LLM/main/L0_MergeRequest_PR pipeline #41933 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

tcherckez-nvidia · 2026-06-08T08:22:36Z

/bot run

tensorrt-cicd · 2026-06-08T08:28:03Z

PR_Github #52715 [ run ] triggered by Bot. Commit: ae181df Link to invocation

tensorrt-cicd · 2026-06-08T12:05:38Z

PR_Github #52715 [ run ] completed with state SUCCESS. Commit: ae181df
/LLM/main/L0_MergeRequest_PR pipeline #41982 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

… 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>

…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>

@tcherckez-nvidia

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>

…lone 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>

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>

tcherckez-nvidia · 2026-06-08T13:06:29Z

/bot run

tensorrt-cicd · 2026-06-08T13:12:30Z

PR_Github #52779 [ run ] triggered by Bot. Commit: d3f0009 Link to invocation

tensorrt-cicd · 2026-06-08T17:36:26Z

PR_Github #52779 [ run ] completed with state SUCCESS. Commit: d3f0009
/LLM/main/L0_MergeRequest_PR pipeline #42035 completed with status: 'SUCCESS'

CI Report

Link to invocation

tensorrt-cicd requested a review from a team as a code owner June 3, 2026 17:47

tensorrt-cicd requested a review from MrGeva June 3, 2026 17:47

tensorrt-cicd assigned tcherckez-nvidia Jun 3, 2026

github-actions Bot assigned tensorrt-cicd Jun 3, 2026

galagam requested review from greg-kwasniewski1 and tcherckez-nvidia June 4, 2026 15:08

tensorrt-cicd force-pushed the repair-bot-bug6248757 branch from 4bd1ee1 to 6fcddfc Compare June 7, 2026 06:03

tcherckez-nvidia changed the title ~~[https://nvbugs/6248757][fix] Keep the optimization enabled~~ [https://nvbugs/6248757][fix] Avoid running all reduce in aux stream Jun 7, 2026

tcherckez-nvidia approved these changes Jun 7, 2026

View reviewed changes

tcherckez-nvidia enabled auto-merge (squash) June 7, 2026 09:09

tensorrt-cicd force-pushed the repair-bot-bug6248757 branch from 1e5386b to a8ffede Compare June 7, 2026 22:40

tensorrt-cicd force-pushed the repair-bot-bug6248757 branch from a8ffede to ae181df Compare June 8, 2026 08:15

tensorrt-cicd added 4 commits June 8, 2026 05:56

tensorrt-cicd force-pushed the repair-bot-bug6248757 branch from ae181df to d3f0009 Compare June 8, 2026 12:56

MrGeva mentioned this pull request Jun 8, 2026

[None][perf] AutoDeploy: fast-path Python-list staging in nest_sequences nv-auto-deploy/TensorRT-LLM#254

Draft

tcherckez-nvidia merged commit 1998324 into NVIDIA:main Jun 8, 2026
7 checks passed

galagam mentioned this pull request Jun 10, 2026

[https://nvbugs/6245279][fix] AutoDeploy: Unwaive accuracy tests #15214

Merged

1 task

Uh oh!

Conversation

tensorrt-cicd commented Jun 3, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram

Estimated code review effort

Possibly related PRs

Suggested reviewers

Uh oh!

suyoggupta commented Jun 5, 2026

Uh oh!

tcherckez-nvidia commented Jun 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tcherckez-nvidia commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 7, 2026

Uh oh!

tcherckez-nvidia commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 7, 2026

Uh oh!

tcherckez-nvidia commented Jun 7, 2026

Uh oh!

tcherckez-nvidia commented Jun 7, 2026

Uh oh!

tcherckez-nvidia commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tcherckez-nvidia commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tcherckez-nvidia commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tensorrt-cicd commented Jun 3, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jun 3, 2026 •

edited

Loading

tcherckez-nvidia commented Jun 7, 2026 •

edited

Loading