[https://nvbugs/6412108][fix] Restore original order — all_reduce the routed partial first, then add the…

[trtllm-agent]

Summary

• Root cause: Sharding-IR refactor swapped the order of expert_output + shared_expert_output and all_reduce, scaling the replicated (unsharded) shared_expert output by world_size (8×) and corrupting MMLU/GSM8K outputs.
• Fix: Restore original order — all_reduce the routed partial first, then add the replicated shared output — and update the comment; also removed the nvbugs/6412108 waiver line.
• 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/6412108

Summary by CodeRabbit

• Bug Fixes

  • Improved the handling of mixture-of-experts outputs to make the final result more stable and consistent across distributed runs.

• Tests

  • Removed a test waiver, so one previously skipped accuracy check will now run normally.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4c98e406-4168-4fda-82ff-43cf04fc61f1

📥 Commits

Reviewing files that changed from the base of the PR and between 3b23a9f and da126a4.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

---

📝 Walkthrough

Walkthrough

The reduction order in Qwen3_5MoeSparseMoeBlock.forward is changed so the routed expert output is all-reduced before adding the shared expert output, rather than summing first then reducing. A test waiver entry for a Qwen3.5 MoE NVFP4 accuracy test is removed.

Changes

MoE reduction fix and waiver update

Layer / File(s)	Summary
Reorder MoE output all-reduce and shared expert addition tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py, tests/integration/test_lists/waives.txt	All-reduce is now applied to the routed expert output first, then the shared expert output is added; the previously waived NVFP4 accuracy test for TestQwen3_5_397B_MoE is unwaived.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15914: Both PRs modify tests/integration/test_lists/waives.txt around Qwen3_5_397B MoE NVFP4 skip/waiver entries.

Suggested reviewers: xinhe-nv, jiaganc, Superjomn

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug fix and matches the main code change.
Description check	✅ Passed	The description explains the root cause, fix, tests, and bug link, so it is mostly complete.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #58496 [ run ] triggered by Bot. Commit: 285807f Link to invocation

[tensorrt-cicd]

PR_Github #58496 [ run ] completed with state SUCCESS. Commit: 285807f
/LLM/main/L0_MergeRequest_PR pipeline #47104 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

[brnguyen2]

The reorder is right for the qwen3.5_moe_35b.yaml path (manual tp_plan with the shared_expert_* entries commented out → shared expert replicated), but it is wrong for the default apply_sharding_hints path, where the model's own hints do shard the shared expert. See the inline comment. Also, the description says a nvbugs/6412108 waiver line was removed — no test-list change is in this diff.

[brnguyen2]

The premise here is wrong for the default pipeline. Qwen3_5MoeMLP does not omit layer_type — lines 643-662 tag all three projections with layer_type="mlp" and tp_mode="colwise"/"colwise"/"rowwise", and its own docstring comment at line 638 still says the opposite of this one ("Tagged layer_type=\"mlp\" so the shared expert is TP-sharded ... lifted to full by the single merge-point all_reduce").

apply_sharding_hints is the default sharding transform (config/default.yaml:134) and applies no shard_layers filter unless a config sets one, so on that path the shared expert's rowwise down_proj emits a per-rank partial. With the new order that partial is added after the only all_reduce and is never reduced — silently 1/world_size of the shared contribution on TP>1, no error. The 8x symptom you fixed comes from qwen3.5_moe_35b.yaml, which disables the hint path (sharding_source: ['manual'], shared_expert entries commented out) and therefore replicates it.

Pick one and make the model self-consistent: either drop the tp_mode/layer_type hints from Qwen3_5MoeMLP so the shared expert is genuinely replicated under both pipelines (then this order is correct and line 638's comment must go), or keep all_reduce(routed + shared) and make the 35b config shard the shared expert. As written, correctness depends on which YAML is loaded.

[trtllm-agent]

NVBug 6412108 is closed as Bug - Fixed. The linked bug appears resolved elsewhere or for a reason that does not prove this PR is redundant. This PR should be judged on its own merits; repair-bot is not auto-closing it.