[None][feat] support non-divisible EP in MoE alltoall and slurm benchmark - #13888
Merged
Conversation
JacobHu-NV
force-pushed
the
feat/non-divisible-ep
branch
2 times, most recently
from
May 9, 2026 09:54
aa060db to
ec776da
Compare
JacobHu-NV
changed the title
JacobHu-NV
force-pushed
the
feat/non-divisible-ep
branch
2 times, most recently
from
May 20, 2026 07:47
40dda9b to
e426525
Compare
JacobHu-NV
marked this pull request as ready for review
May 20, 2026 07:54
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49391 [ run ] triggered by Bot. Commit: |
Contributor
📝 WalkthroughWalkthroughThis PR extends MoE expert parallelism to support non-divisible expert distributions via ceil/floor contiguous partitioning and introduces SLURM compact-packing mode for efficient cross-node GPU scheduling. CUDA kernels, PyTorch interfaces, and tests are updated to handle uneven expert-to-rank assignments; SLURM job submission adds optional per-worker hostfiles and GPU mapping. ChangesNon-Divisible Expert Parallel Support
SLURM Compact Packing Scheduling
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tests/unittest/_torch/modules/moe/test_moe_comm.py (1)
450-456: ⚡ Quick winKeep the feasibility gate aligned with the production whitelist.
CommunicationFactorynow allowsAllGatherReduceScatterfor non-divisible EP, but this helper still skips every backend exceptNVLinkOneSided. That makes the fallback path impossible to cover in this suite.Suggested fix
- if config.num_experts % config.ep_size != 0 and comm_type != COMM_NVLINK_ONE_SIDED: - # Only NVLinkOneSided supports non-divisible EP (ceil/floor partitioning). - # Other comm types still require num_experts divisible by ep_size. + if config.num_experts % config.ep_size != 0 and comm_type not in ( + COMM_NVLINK_ONE_SIDED, + COMM_ALLGATHER_RS, + ): + # Non-divisible EP is currently covered for NVLinkOneSided and the + # AllGatherReduceScatter fallback path. return ( f"comm_type={comm_type} requires num_experts divisible by ep_size, " f"got num_experts={config.num_experts}, ep_size={config.ep_size}" )QA list updates look unnecessary here because this remains unittest-only coverage. As per coding guidelines, "Coverage expectations: ... Note missing negative tests ..." and "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."
🤖 Prompt for 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. In `@tests/unittest/_torch/modules/moe/test_moe_comm.py` around lines 450 - 456, The feasibility gate incorrectly only permits non-divisible EP for COMM_NVLINK_ONE_SIDED while CommunicationFactory now also allows COMM_ALLGATHER_REDUCESCATTER; update the conditional in the helper (the block that checks config.num_experts % config.ep_size != 0 and comm_type != COMM_NVLINK_ONE_SIDED) to also accept COMM_ALLGATHER_REDUCESCATTER (e.g., change the single-value exclusion to a membership test or OR-check), and adjust the returned error message text to reflect which comm types require num_experts divisible by ep_size so tests can exercise the fallback path; locate references to COMM_NVLINK_ONE_SIDED and COMM_ALLGATHER_REDUCESCATTER in this file to apply the change.
🤖 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 `@examples/disaggregated/slurm/benchmark/submit.py`:
- Line 713: Remove the erroneous f-string prefix on the literal
f"--no-container-mount-home --mpi=pmix --overlap -N 1 -n 1" in submit.py (this
is the list/argument element used when building the sbatch/submit command);
replace it with a plain string "--no-container-mount-home --mpi=pmix --overlap
-N 1 -n 1" so the linter F541 is not triggered.
- Line 441: The assignment to compact_packing is using
bool(hw_config.get('compact_packing', False)) which will incorrectly coerce
string values like "false" to True; change the logic in the submit code that
reads hw_config.get('compact_packing') (the compact_packing variable) to accept
actual booleans and parse string inputs explicitly (e.g., treat case-insensitive
"true"/"1"/"yes" as True and "false"/"0"/"no" as False) or fallback to the
default False when the key is missing or unrecognized; ensure the code first
checks isinstance(value, bool) and otherwise normalizes str values before
setting compact_packing.
In `@tensorrt_llm/_torch/modules/fused_moe/interface.py`:
- Around line 449-458: The current branch builds initial_global_assignments from
moe_load_balancer_config.num_local_slots even when the real load balancer is
missing, causing mismatched initial_local_expert_ids; change the logic in the
constructor/code that sets initial_global_assignments (where
moe_load_balancer_config, init_expert_size_per_partition and
initial_global_assignments are computed) to detect the actual availability of
the load balancer (i.e., only use moe_load_balancer_config.num_local_slots when
get_moe_load_balancer() returned a valid balancer object); if the balancer is
absent, set initial_global_assignments to the safe sequential mapping
list(range(self.num_experts)) instead; apply the same defensive check and
fallback in the other identical block (the second occurrence around
initial_local_expert_ids).
---
Nitpick comments:
In `@tests/unittest/_torch/modules/moe/test_moe_comm.py`:
- Around line 450-456: The feasibility gate incorrectly only permits
non-divisible EP for COMM_NVLINK_ONE_SIDED while CommunicationFactory now also
allows COMM_ALLGATHER_REDUCESCATTER; update the conditional in the helper (the
block that checks config.num_experts % config.ep_size != 0 and comm_type !=
COMM_NVLINK_ONE_SIDED) to also accept COMM_ALLGATHER_REDUCESCATTER (e.g., change
the single-value exclusion to a membership test or OR-check), and adjust the
returned error message text to reflect which comm types require num_experts
divisible by ep_size so tests can exercise the fallback path; locate references
to COMM_NVLINK_ONE_SIDED and COMM_ALLGATHER_REDUCESCATTER in this file to apply
the change.
🪄 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: 778f0924-4f42-4cab-8f10-6e20b58a081f
📒 Files selected for processing (11)
cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cucpp/tensorrt_llm/thop/moeAlltoAllOp.cppexamples/disaggregated/slurm/benchmark/README.mdexamples/disaggregated/slurm/benchmark/config.yamlexamples/disaggregated/slurm/benchmark/disaggr_torch.slurmexamples/disaggregated/slurm/benchmark/run_benchmark.shexamples/disaggregated/slurm/benchmark/start_worker.shexamples/disaggregated/slurm/benchmark/submit.pytensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.pytensorrt_llm/_torch/modules/fused_moe/interface.pytests/unittest/_torch/modules/moe/test_moe_comm.py
Collaborator
|
PR_Github #49391 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49538 [ run ] triggered by Bot. Commit: |
qiaoxj07
reviewed
May 21, 2026
Collaborator
|
PR_Github #49538 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49813 [ run ] triggered by Bot. Commit: |
tianyuz-nv
added a commit
to tianyuz-nv/TensorRT-LLM
that referenced
this pull request
May 22, 2026
Port the per-rank hostfile + gpu_map launcher mechanism from PR NVIDIA#13888 (Jianbo Hu, "support non-divisible EP in MoE alltoall and slurm benchmark") into the in-tree DWDP examples launcher trio (submit_dwdp.py, disaggr_torch_dwdp.slurm, start_worker_dwdp.sh). Motivation: SLURM block distribution requires (n_ctx + n_gen * gen_tp) to be a multiple of gpus_per_node, otherwise GEN tensor-parallel ranks get split across nodes/trays and incur a large allreduce penalty. Previously worked around by adding empty CTX slots (over-provisioning) or picking dwdp_group multiples that happen to divide cleanly — both fragile and incompatible with Mode B non-uniform expert ranges where num_ctx is e.g. 5 or 6. Dual-path design: * Divisible case (num_ctx_gpus % gpus_per_node == 0): Use the legacy --nodelist + -N + --ntasks-per-node srun command. Block distribution gives the natural rank-to-node mapping; trtllm-serve picks its CUDA device from SLURM_LOCALID. No hostfile / gpu_map files are emitted. Perf-optimal path. * Non-divisible case (Mode B dwdp3 dg=2, dwdp5 dg=1, etc.): Emit hostfile_mpi_worker_base.txt and gpu_map_mpi_worker_base.txt under log_dir, iterating CTX servers then GEN servers so global rank ordering matches split_world_comm's ctx_cfgs + gen_cfgs. disaggr_torch_dwdp.slurm rewrites <nodeN_placeholder> tokens to real hostnames at runtime. srun uses SLURM_HOSTFILE + --distribution=arbitrary to pin every rank's host placement. DWDP-specific deviation from PR NVIDIA#13888: start_worker_dwdp.sh does NOT export CUDA_VISIBLE_DEVICES. DWDP relies on intra-node peer GPU discovery (VA composite cuMemMap of peer MNNVL fabric handles, UCX cuda_ipc / cuda_copy intra-node KV transports, PyTorch peer device enumeration); restricting CUDA visibility to a single GPU breaks these paths and causes a 15% per-CTX-GPU regression (TPOT unchanged, TTFT std blows up 3x). With our allocate_gpus's sequential cursor, gpu_id == SLURM_LOCALID for every rank, so trtllm-serve's internal LOCALID-based device selection already lands each rank on the correct GPU. The gpu_map file is kept for diagnostics and audit logging only. Empirical per-CTX-GPU req/s (DSv3-FP4-V2.1, ISL=8192/OSL=1, conc=256): dwdp4 dg=1 (8 GPU divisible) 3.46 dwdp3 dg=4 (16 GPU divisible) 3.41 dwdp5 dg=4 (24 GPU divisible) ~3.3 (extrapolated from R1b) dwdp3 dg=2 (10 GPU non-divisible) 3.51 dwdp5 dg=1 (9 GPU non-divisible) 3.42 TPOT 14-18 ms across all configs (kernel-neutral). Enables the non-uniform / Mode B perf configurations used in the follow-up dwdp_size=3/5 experiments in this refactor. Co-authored-by: Jianbo Hu <jacohu@nvidia.com> Signed-off-by: tianyuz-nv <tianyuz@nvidia.com>
- submit.py: fix calculate_nodes to match assign_server_round_robin's per-worker whole-node ownership semantic. Pooled formula ceil(world_size * num_servers / gpus_per_node) under-counts nodes when world_size is not a multiple of gpus_per_node (qiaoxj07 review, e.g. 2 ctx workers x ws=6 + 1 gen worker x ws=4 with gpus_per_node=4 -> old formula gave 4 but allocator needed 5 -> IndexError). - submit.py: drop erroneous f-string prefix on line introduced by e2157c3 (ruff F541, CodeRabbit review). - test_moe_comm.py: relax check_feasibility to also allow AllGatherReduceScatter under non-divisible EP, matching the production fallback in CommunicationFactory. Signed-off-by: JacobHu-NV <266902545+JacobHu-NV@users.noreply.github.com>
Match the formatting the project's yapf hook expects (one element per line + dangling comma + closing paren on its own line), as enforced by the PR check pre-commit on 221c9d6. Signed-off-by: JacobHu-NV <266902545+JacobHu-NV@users.noreply.github.com>
…divisible fast path The ceil/floor expert partitioning introduced in commit f89fc92 added a modest dispatch-kernel overhead in the divisible case (num_experts % ep_size == 0) at small batch sizes, because compute_target_rank_id was performing 2 integer divisions + 1 multiply + 1 branch per TOP_K call when NVCC failed to CSE the per-call base/remainder recompute out of the per-token loop. This change: 1. Lifts the base/remainder computation into the kernel body once, just before the TOP_K loop, and passes the precomputed values into compute_target_rank_id. NVCC no longer needs to discover the loop invariant. 2. Adds a uniform-partition fast path: when remainder == 0, compute_target_rank_id returns `expert_id / base` directly, matching the pre-PR single-divide cost exactly. Marks compute_target_rank_id __forceinline__ so the call collapses fully in the dispatch kernel. Benchmark on GB200 NVL72, DeepSeek-V3 profile (hidden=7168, top_k=8, num_experts=256, BF16, NVLinkOneSided, 200 iters / 20 warmup): ep=8 dispatch divisible-case overhead vs baseline (PR_v3 = this change): batch BL µs PR µs Δ% 1 25.61 25.66 +0.21 4 26.42 26.57 +0.55 16 26.96 27.04 +0.30 64 31.70 31.80 +0.31 256 59.73 59.99 +0.44 1024 173.87 174.10 +0.13 2048 322.12 322.29 +0.05 ep=16 dispatch: similar; combine kernel untouched and Δ stays in ±0.5% run-to-run noise across both ep sizes. Worst-case Δ across 24 measured (ep, batch) points is +1.93% (well below ~6% rank-stdev at small batches). Non-divisible case (ep=7, num_experts=256, base=36, remainder=4) reaches 735 GB/s dispatch / 587 GB/s combine at b=2048, the same ~80% NVLink-peak regime as ep=8 in the divisible case. Signed-off-by: JacobHu-NV <266902545+JacobHu-NV@users.noreply.github.com>
Only the NVLinkOneSided dispatch kernel + the CuteDsl MoE backend have
been validated for the ceil/floor expert partition introduced earlier in
this PR. Other MoE backends (CutlassFusedMoE, TRTLLMGenFusedMoE,
DeepGemmFusedMoE, DenseGEMMFusedMoE, TritonFusedMoE, MegaMoEDeepGemm) and
VanillaMoE still assume uniform partitioning and would silently produce
wrong local-expert ranges if instantiated with num_experts % ep_size != 0.
Add an explicit class-attribute opt-in flag and check it in MoE.__init__
right after self.ep_size is set:
- MoE base class: _supports_non_divisible_ep = False (default reject)
- CuteDslFusedMoE: _supports_non_divisible_ep = True (only opt-in)
- ConfigurableMoE: _supports_non_divisible_ep = True (wrapper passes;
inner backend's
own gate is the
authoritative
check)
- VanillaMoE: explicit ValueError (does not inherit MoE,
cannot share the base gate)
Behavior matrix for ConfigurableMoE wrapping a concrete backend:
ConfigurableMoE -> CuteDslFusedMoE + non-divisible -> allowed (both pass)
ConfigurableMoE -> CutlassFusedMoE + non-divisible -> blocked at inner init
ConfigurableMoE -> TRTLLMGen / etc. + non-divisible -> blocked at inner init
Any backend + divisible -> unchanged
Validated end-to-end on Kimi-K2-Thinking-NVFP4 (num_experts=384) over a
running disaggregated benchmark stack:
- ep=5 (384%5=4) ConfigurableMoE->CuteDsl : ctx + gen workers ready
- ep=6 (384%6=0) ConfigurableMoE->CuteDsl : ctx + gen workers ready
- ep=7 (384%7=6) ConfigurableMoE->CuteDsl : ctx + gen workers ready
- ep=5 ConfigurableMoE->Cutlass (forced) : blocked with clear ValueError
from CutlassFusedMoE.__init__
Also exercises the existing tests/_torch/modules/moe/test_moe_comm.py
matrix end-to-end on the patched build (NVLinkOneSided), all 18 cases pass
(6 non-divisible + 12 divisible) including verify_dispatch_results and
verify_combine_results (rtol=0.02 atol=0.15, fp8 low-precision combine
included).
Signed-off-by: JacobHu-NV <266902545+JacobHu-NV@users.noreply.github.com>
JacobHu-NV
force-pushed
the
feat/non-divisible-ep
branch
from
May 26, 2026 07:37
db7dd51 to
f425a5e
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50338 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50310 [ run ] completed with state |
Collaborator
|
PR_Github #50338 [ run ] completed with state
|
xxi-nv
self-requested a review
May 26, 2026 23:48
Collaborator
|
@JacobHu-NV Could you help to clarify whether the case below is supported? If not, do we have any guard to raise error? |
The ceil/floor partition silently produces zero-expert ranks when num_experts < ep_size (e.g. num_experts=2, ep_size=8 leaves ranks 2-7 with empty partitions). No MoE backend or comm strategy supports this end-to-end: the moeAlltoAllOp has a backstop check, but the AllGatherReduceScatter fallback bypasses it, leading to cryptic downstream shape errors. Raise a clear ValueError upfront. Signed-off-by: JacobHu-NV <266902545+JacobHu-NV@users.noreply.github.com>
Collaborator
Author
Good point, I added a new commit to fix this |
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50425 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50425 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #50579 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #50579 [ run ] completed with state |
zongfeijing
requested review from
a team
and removed request for
FrankD412,
QiJune,
bo-nv,
kaiyux and
reasonsolo
May 28, 2026 01:25
qiaoxj07
approved these changes
May 28, 2026
bmarimuthu-nv
pushed a commit
to nv-auto-deploy/TensorRT-LLM
that referenced
this pull request
May 28, 2026
…mark (NVIDIA#13888) Signed-off-by: JacobHu-NV <266902545+JacobHu-NV@users.noreply.github.com> Signed-off-by: JacobHu-NV <jacohu@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request enhances support for non-divisible expert parallelism (EP) in Mixture-of-Experts (MoE) communication by implementing ceil/floor partitioning for expert assignment across ranks. The changes ensure that MoE communication and computation remain correct and efficient when the number of experts is not evenly divisible by the number of ranks. This includes updates to CUDA kernels, Python interfaces, strategy selection logic, and comprehensive unit tests.
Key changes include:
Core algorithm and kernel updates:
moeAlltoAllKernels.cu) to use ceil/floor partitioning incompute_target_rank_id, allowing non-divisible expert-to-rank mapping and ensuring all experts are correctly assigned even whennum_experts % ep_size != 0. [1] [2]num_expertsto be divisible byep_sizein the C++ operator, reflecting the new partitioning logic.Python interface and strategy selection:
_compute_ep_partitionin the Python interface to mirror the kernel's partitioning logic, ensuring consistent expert slot assignment and correct initialization in both standard and load-balanced MoE scenarios. [1] [2] [3] [4]AllGatherReduceScatterwhen non-divisible EP is detected, and enforced that only whitelisted methods (NVLinkOneSided, AllGather) are allowed for non-divisible configurations. [1] [2]Testing and validation:
These changes collectively make the MoE implementation robust to non-uniform expert distributions, improving flexibility and correctness in distributed training scenarios.
Performance
Setup: GB200 NVL72 (4-node alloc, 4 GPUs/node), DeepSeek-V3 profile (
hidden=7168,top_k=8,num_experts=256), BF16, NVLinkOneSided backend.tests/microbenchmarks/bench_moe_comm.py --backend NVLINK_ONE_SIDED --profile deepseek_v3 --perfect_router --kernel_breakdown --iter_stats -b 1 -e 2048 -f 2. All values are mean across MPI ranks of per-rank mean latency over 200 iterations (20 warmup).Dispatch latency (µs)
PR shows a tiny dispatch-side overhead vs baseline (worst case +1.93% at ep=16 b=2; most points <+1%, large batches ≤+0.1%). This is the cost of the new ceil/floor routing logic in
compute_target_rank_id(extraimul+ branch even on the divisible fast path); at small batch the worst Δ% (≤2%) is well within the ~5-6% iteration-to-iteration stdev observed per rank (and ~2-3% inter-rank stdev), and at large batch it amortizes to ≤±0.5% (run-to-run noise).Combine latency (µs)
(Combine kernel is unchanged by this PR; table included as a noise-floor reference.)
Non-divisible EP:
ep_size=7(256/7 → base=36, remainder=4)PR-only (the pre-PR kernel
TORCH_CHECKsnum_experts % ep_size == 0and refuses to launch).Summary by CodeRabbit
New Features
compact_packingoption for GPU worker placement in distributed SLURM deploymentsImprovements
Documentation
compact_packingusage guidance