[https://nvbugs/6412108][fix] AutoDeploy:Shard shared Qwen3.5 experts

[greg-kwasniewski1]

Summary

Fix the Qwen3.5-397B MoE NVFP4 accuracy regression (nvbugs/6412108) by sharding the shared expert instead of keeping it replicated.

Root cause

The shared expert (Qwen3_5MoeMLP) had its torch_linear_simple ops left untagged, and qwen3.5_moe_400b.yaml set apply_sharding_hints.shard_layers: ["moe","delta","mha"], which excludes untagged ("unknown") nodes. So the shared expert stayed replicated (full output on every rank) while the routed experts were sharded. The single merge-point expert_output + shared_expert_output followed by all_reduce then scaled the replicated shared output by world_size (8x), dropping MMLU from ~85 to ~0.05.

Fix

Shard the shared expert like the routed experts rather than special-casing the reduction order:

• Tag the shared expert's gate/up (colwise) and down (rowwise) linears with layer_type="mlp" in modeling_qwen3_5_moe.py.
• Remove the shard_layers whitelist from qwen3.5_moe_400b.yaml so all tagged nodes shard.
• Keep the existing add-then-all_reduce merge point: both routed and shared branches now emit per-rank partials, so one reduction on the sum lifts both to full.

This decouples modeling-code numerical correctness from an optional YAML knob (the whitelist).

Relation to #15922

This supersedes the alternative fix in #15922, which instead reordered the collective (all_reduce the routed partial first, then add the replicated shared output). That approach keeps the shared expert replicated and leaves correctness dependent on the yaml whitelist. Local benchmarking (Qwen3.5-397B-A17B-NVFP4, 8xB200, ISL/OSL 1024, concurrency 1/16/32/64) shows replicating vs sharding the shared expert is within ~1-2% throughput (run-to-run noise, no consistent winner), so there is no perf reason to keep it replicated; sharding both paths is the simpler, self-consistent fix.

Test plan

• TestQwen3_5_397B_MoE::test_nvfp4[8] accuracy: MMLU ~86.8 (was ~0.05), passes threshold; unwaived here.
• test_sharding_num_correctness.py for modeling_qwen3_5_moe across tp-only / ep-only / tep / attn-dp (normal pass + SHARDING_IR_SABOTAGE=1 true-negative).
• tests/unittest/auto_deploy/multigpu/ suite.
• CI green (triggered after open).

Risk

library-visible (AutoDeploy sharding / modeling).

Summary by CodeRabbit

• Performance
  • Improved deployment configuration for Qwen3.5 MoE models with enhanced sharding and multi-stream execution.
  • Optimized expert-layer parallelism to support more efficient inference.
• Compatibility
  • Improved coordination of model outputs during distributed execution.
• Testing
  • Expanded integration test coverage by removing an outdated waiver for Qwen3.5 MoE accuracy testing.

[greg-kwasniewski1]

/bot run

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Qwen3.5 MoE sharding

Layer / File(s)	Summary
Sharding and streaming configuration examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml	Replaces shard_layers with an lm_head filter and enables compile-stage MoE/GEMM multi-streaming plus logits gathering.
MLP sharding annotations and waiver tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py, tests/integration/test_lists/waives.txt	Tags shared-expert MLP projections for sharding and removes the related NVFP4 test waiver.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: yechank-nvidia, Hudayday, litaotju, aswinvisva, lfr-0531

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title matches the required ticket/type format and clearly summarizes the shared-expert sharding fix.
Description check	✅ Passed	The description is complete and well-structured, covering summary, root cause, fix, related work, tests, and risk.

✨ 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

[greg-kwasniewski1]

/bot run

[tensorrt-cicd]

PR_Github #58811 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

[tensorrt-cicd]

PR_Github #58811 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

[tensorrt-cicd]

PR_Github #58813 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

[tensorrt-cicd]

PR_Github #58813 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

[tensorrt-cicd]

PR_Github #58815 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

[tensorrt-cicd]

PR_Github #58815 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance.

Link to invocation

[chzblych]

/bot run

[tensorrt-cicd]

PR_Github #58830 [ run ] triggered by Bot. Commit: 9f8d009 Link to invocation

[tensorrt-cicd]

PR_Github #58830 [ run ] completed with state SUCCESS. Commit: 9f8d009
/LLM/main/L0_MergeRequest_PR pipeline #47381 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

[greg-kwasniewski1]

/bot run

[greg-kwasniewski1]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59159 [ run ] triggered by Bot. Commit: de481d3 Link to invocation