[None][fix] Keep chunked-MoE size vectors identical across attention-DP ranks - #16745
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThe MoE scheduler now patches token-count entries for every DP rank with an empty chunk, using chunk-0 values for consistent variable-size collective shape vectors. ChangesDP empty-chunk synchronization
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py`:
- Around line 638-647: Restrict the empty-chunk substitution loop over
all_rank_num_tokens_list to the varsize collective path by guarding it with not
moe.enable_alltoall. Preserve the earlier AllToAll zero-to-one sizing
established in the surrounding scheduler logic, and leave
all_rank_chunk_size_list unchanged as the empty-chunk reference.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 205070eb-4946-4acf-96d1-24510eacc946
📒 Files selected for processing (1)
tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py
|
/bot run |
|
PR_Github #61030 [ run ] triggered by Bot. Commit: |
|
PR_Github #61030 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61037 [ run ] triggered by Bot. Commit: |
fe63115 to
20444da
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61045 [ run ] triggered by Bot. Commit: |
|
PR_Github #61037 [ run ] completed with state |
|
PR_Github #61045 [ run ] completed with state
|
20444da to
00c8ffa
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61131 [ run ] triggered by Bot. Commit: |
|
PR_Github #61131 [ run ] completed with state
|
00c8ffa to
0700d8a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61497 [ run ] triggered by Bot. Commit: |
|
PR_Github #61497 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61602 [ run ] triggered by Bot. Commit: |
|
PR_Github #61602 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61651 [ run ] triggered by Bot. Commit: |
|
PR_Github #61651 [ run ] completed with state
|
…DP ranks On cross-node IB (no MNNVL) the MoE comm lane is DeepEP; when a heavy prefill exceeds moe_config.max_num_tokens the forward is chunked, and DeepEP cannot run chunked, so each chunk goes through the varsize allgather/reducescatter fallback. The collectives cannot take an empty chunk, so a rank whose chunk is empty resends its chunk 0 instead -- but the size-vector update only covered the local rank's entry. Ranks then post mismatched counts for the same collective and NCCL deadlocks: all GPUs spin at 100%, no error, watchdog abort after 300 s. Fix: patch every rank's empty entries from the shared chunk-size table, so all ranks derive identical size vectors. No-op for the alltoall path (its existing 0->1 patch already yields the same values). Repro: Qwen3.5-397B-A17B-NVFP4, 2-node x86 B200, DEP16 (attention-DP, moe_ep=16, moe_tp=1), moe_config.max_num_tokens=32768, 32k-ISL bench at concurrency 1 -- deterministic hang at iter 2 without this fix, clean 208k+ iterations with it. Signed-off-by: dongfengy <99041270+dongfengy@users.noreply.github.com> Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
0700d8a to
9089e19
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #61954 [ run ] triggered by Bot. Commit: |
|
PR_Github #61954 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62003 [ run ] triggered by Bot. Commit: |
|
PR_Github #62003 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62038 [ run ] triggered by Bot. Commit: |
|
PR_Github #62038 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62138 [ run ] triggered by Bot. Commit: |
|
PR_Github #62138 [ run ] completed with state |
4 participants
Description
Under attention-DP, when a heavy prefill exceeds
moe_config.max_num_tokensthe MoE forward is chunked. On the varsize allgather/reduce-scatter comm lane (the fallback whenever the alltoall lane cannot run chunked, e.g. DeepEP on multi-node IB), the collectives cannot take an empty chunk, so a rank whose chunk is empty re-sends its chunk-0 data instead — but the size-vector update in_forward_multiple_chunksonly patched the local rank's entry ofall_rank_num_tokens_list. Different ranks then derive different size vectors for the same collective, and NCCL deadlocks: all GPUs spin at 100%, no error until the watchdog aborts (~300 s).The fix patches every rank's empty entries from
all_rank_chunk_size_list(the shared ground truth all ranks compute identically), so the per-chunk size vectors are identical across ranks. This is a no-op for the alltoall lanes, whose existing 0→1 substitution already yields identical vectors.Test Coverage
Reproduced deterministically with Qwen3.5-397B-A17B (NVFP4) wide-EP EP16 (attention-DP) across 2×8 B200 over InfiniBand, 32k-ISL prefills with
moe_config.max_num_tokens=32768: the run deadlocks at the second chunked iteration without this patch and completes 210k+ iterations cleanly with it (same wheel A/B). Accuracy sanity of the same setup is unaffected (single-chunk paths untouched).Trigger conditions covered: attention-DP + chunked MoE (cap exceeded) + at least one rank with fewer tokens than chunks (mixed prefill/decode across DP ranks — the steady state of any real serving workload).
PR Checklist
🤖 Generated with Claude Code
Dev Engineer Review
tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py(ExternalCommMoEScheduler._forward_multiple_chunks), expanded the DP empty-chunk substitution so that for every chunk indexidx_chunkand every DP rankj:all_rank_chunk_size_list[j][idx_chunk] == 0, the code patchesall_rank_num_tokens_list[idx_chunk][j]toall_rank_chunk_size_list[j][0].x_list/router_logits_list/input_ids_listfrom empty chunks to chunk-0 remains in place; the change specifically updates the DP collective shape inputs (all_rank_num_tokens_list) to eliminate rank divergence.QA Engineer Review
No test changes.