Repository navigation
[EPLB] Add contention-aware expert migration batching - #52641
Merged
robertgshaw2-redhat merged 44 commits intoOct 3, 2026
Merged
robertgshaw2-redhat merged 44 commits into
robertgshaw2-redhat merged 44 commits into
Conversation
DOCCA0
marked this pull request as ready for review
August 17, 2026 23:56
DOCCA0
requested review from
ProExpertProg,
WoosukKwon,
houseroad,
mgoin,
robertgshaw2-redhat,
tlrmchlsmth,
yewentao256 and
youkaichao
as code owners
August 17, 2026 23:56
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
DOCCA0
force-pushed
the
agent/eplb-migration-batching
branch
from
August 17, 2026 23:59
23bd32d to
394576d
Compare
DOCCA0
force-pushed
the
agent/eplb-migration-batching
branch
from
August 28, 2026 11:40
394576d to
0890d07
Compare
ilmarkov
reviewed
Aug 31, 2026
ilmarkov
left a comment
ilmarkov
left a comment
Contributor
There was a problem hiding this comment.
The PR looks good. Some cosmetic changes required.
Regarding the benchmarking, why didn't you use the default nixl eplb communicator? Also, from the command line it's unclear what were the eplb parameters (most importantly step interval).
Contributor
Author
I did both gloo and nixl. But the bench.json file of nixl is coverd by bench.json of torch_gloo . |
SageMoore
reviewed
Aug 31, 2026
SageMoore
left a comment
SageMoore
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the work! I have some initial comments/questions.
DOCCA0
force-pushed
the
agent/eplb-migration-batching
branch
from
September 1, 2026 04:41
0890d07 to
79c6ea4
Compare
ilmarkov
reviewed
Sep 1, 2026
ilmarkov
left a comment
ilmarkov
left a comment
Contributor
There was a problem hiding this comment.
Thank you for the benchmarks updates.
- Short comment on
set_transfer_contextin_execute_migration_batch - I would suggest to make the limit of peer number in a batch configurable (at the moment each rank has only one peer per batch which could be left as default value). It could reduce the number of steps, and allow saturate high-bandwidth NICs.
Contributor
|
/ci run |
|
❌ @ilmarkov, Only reviewers with write access can use CI commands before CI is delegated to the PR author. |
Member
|
/ci run |
|
✅ Triggered Buildkite CI #87080 for commit |
DOCCA0
force-pushed
the
agent/eplb-migration-batching
branch
2 times, most recently
from
September 3, 2026 15:55
780ae15 to
335baed
Compare
Contributor
|
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: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. WalkthroughEPLB expert migrations now support deterministic, rank-disjoint batching. Configuration enables batching by default, and transfer call sites forward the setting. Tests cover scheduling, buffered transfers, configuration, and validation behavior. ChangesEPLB migration batching
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant transfer_layer
participant move_to_buffer
participant migration_scheduler
participant EplbCommunicator
participant move_from_buffer
transfer_layer->>move_to_buffer: pass migration indices and batching flag
move_to_buffer->>migration_scheduler: build and schedule migration batches
migration_scheduler-->>move_to_buffer: return rank-disjoint batches
move_to_buffer->>EplbCommunicator: set context and post sends or receives
EplbCommunicator-->>move_to_buffer: execute each batch
move_to_buffer-->>move_from_buffer: return TransferMetadata
move_from_buffer->>move_from_buffer: copy buffered weights into expert weights
Merge Risk: ⚪ Minimal · up to This change enables contention-aware EPLB expert-migration batching by default while retaining an opt-out path. The supplied test and implementation summaries identify no unresolved correctness, availability, or deployment risk that blocks merging. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
❌ @DOCCA0, A reviewer with write access must run |
Member
|
/ci run |
|
✅ Triggered Buildkite CI #87146 for commit |
DOCCA0
commented
Sep 4, 2026
DOCCA0
left a comment
DOCCA0
left a comment
Contributor
Author
There was a problem hiding this comment.
reviewed line by line
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Signed-off-by: DOCCA0 <2947323341@qq.com>
robertgshaw2-redhat
added
the
ready
Collaborator
|
/ci run |
|
✅ @DOCCA0, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #92422 for commit |
robertgshaw2-redhat
approved these changes
Oct 1, 2026
robertgshaw2-redhat
left a comment
robertgshaw2-redhat
left a comment
Collaborator
There was a problem hiding this comment.
approving given Sage/Ilia own this area
robertgshaw2-redhat
enabled auto-merge (squash)
October 1, 2026 14:55
Collaborator
|
/ci run |
|
✅ Triggered Buildkite CI #92679 for commit |
Collaborator
CI selector (shadow): 171 test steps (243 jobs) instead of 71 (91 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (171)
Would skip (today's rules run them) (4)
Would add (today's rules do not run them) (104)
AMD mirrors: would skip (2)
AMD mirrors: would add (76)
6 changed files · base |
Contributor
Author
|
/ci retry --allow-stale |
Contributor
Author
|
/ci retry |
|
✅ Triggered Buildkite CI #92715 for commit |
Contributor
Author
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #92715. |
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
7 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.
Partially addresses #31671.
Summary
During async EPLB, several expert weights can be transferred through the same rank and NIC at once. These traffic bursts compete with inference and can increase serving latency.
This change adds an opt-in scheduler that groups migrations by source and destination rank, then splits them into batches. In each batch, a rank communicates with at most one peer, while independent rank pairs still transfer concurrently. Expert weights still move one layer at a time, but schedules for all MoE layers are prepared once per rebalance cycle to reduce Python overhead.
enable_migration_batchingremainsfalseby default. Setting it totrueenables the new path only whenuse_async=true; synchronous EPLB always keeps the existing path because inference is paused and serial batches would extend the pause. NIXL transfer context is restored before every batch becauseexecute()drains it.I checked open EPLB migration/scheduling PRs and found no duplicate implementation. AI assistance was used; I (a human contributor) reviewed every changed line and am responsible for the contribution.
Tests
Result:
12 passed; selected pre-commit hooks passed.End-to-end serving benchmark
One four-node run used Ubuntu 22.04.5, one Quadro RTX 6000 24 GB GPU per node, Ray 2.56.1, NIXL 1.3.2 over 10 GbE without RDMA, and prefix caching enabled. The model was
Qwen/Qwen3-30B-A3B-Instruct-2507with TP=4 and 32-way concurrency. Each case used a fresh server; 500/500 requests completed.Server command:
Bench command:
Async batching improved request throughput by 2.28%, TTFT P50/P99 by 1.18%/1.30%, TPOT P50/P99 by 1.94%/0.87%, and E2EL P50/P99 by 3.18%/1.50%.
The head node's total network traffic (bytes received + bytes transmitted) was sampled from Linux byte counters once per second; it is not a
vllm bench servemetric. NIC P99 decreased from 599.81 to 538.40 MB/s (10.24%). The plots and unmodified JSON/log files are inresults/serving_nixl_20260911_async_random500.Scheduler profiling
The profiling run used the same server and workload parameters, with only NVTX-only Nsight capture enabled. The exact commands are documented in
COMMANDS.md. Open the fourrank*.nsys-repfiles with Nsight Systems 2024.5.1 or newer and search foreplb: schedule migration batches.Each range is one scheduler call covering all 48 MoE layers in a rebalance cycle, rather than one call per layer. The transfers themselves remain layer-by-layer. The measured per-rank results were:
The four ranks schedule concurrently, so their times are not added. Across the 34 cycle-wide scheduler calls, rank 0 had the largest measured total scheduling time: 302.081 ms. The independent clean A/B benchmark completed 105.932 s faster with batching. Comparing these two measured values gives a scheduler-cost-to-serving-time-saved ratio of 0.285%.
Thanks to my advisor Massimo Tornatore (massimo.tornatore@polimi.it) and Assistant Professor Qiaolun Zhang (qiaolun.zhang@polimi.it) for their professional support.