[kpham-sgl]

Motivation

An EPLB rebalance stops serving for as long as it runs: 1.44-2.15 s per rebalance on a DEP2 DeepSeek-V4.1-Flash decode (GB300, MegaMoE), and 3.7-4.5 s per rebalance on the MI355X prefill EP8 measured in #43754. The rebalance generator is advanced at the end of each forward pass on the thread that launches forward passes. It computes the placement in Python there, builds one P2P op per tensor per moved expert there, submits the transfers on the default communicator, whose NCCL stream waits for the current forward and whose completion the next forward waits for, and copies the staged experts into the live weights on the compute stream. The default single-chunk mode also wraps all of it in two synchronize() calls for timing.

--eplb-rebalance-layers-per-chunk spreads the per-layer work over passes but changes none of the things that stall a step: the placement is computed whole before the first chunk; each chunk's Python stays on the forward thread; each chunk's transfers still sit between two forwards on the compute stream; the staging-to-weight copies must sit between two forwards. The total GPU time is unchanged, split into 40/N pauses.

A first version of this PR moved the transfers to NCCL point-to-point ops on a side stream. On a DEP16 MegaMoE decode (16 ranks, 4 nodes) it hung within two seconds of the first rebalance and every rank died in DeepGEMM's Grid sync timeout: on SM100 MegaMoE is a whole-grid clustered launch that needs every CTA resident and is given all SMs but a margin of 2 (mega_moe.py, SGLANG_OPT_DEEPGEMM_MEGA_MOE_RESERVED_SMS), so a rank whose NCCL kernel holds SMs cannot bring its MoE grid up, while the peer that would complete the NCCL transfer sits inside its own MoE kernel waiting at the cross-rank barrier for this rank. Any transfer that runs as a GPU kernel concurrently with that MoE kernel can deadlock it. This version moves the weights with no GPU kernel at all.

Modifications

• New eplb_async_migrator.py. Each rank pulls the experts it needs from its peers' live weights through the Mooncake transfer engine with one-sided reads (batch_transfer_sync_read, the primitive the elastic-EP expert backup and the remote-instance weight loader already use), into one layer of staging buffers; DMA on the NIC and NVLink, no SMs, no sender participation, no pairing of sends and receives. plan_layer_pull is the pure planner: unchanged slot, same-GPU copy, free-rider on an earlier local slot, otherwise a read from a holder on the same node if there is one, spread across holders. The reads run on a worker thread; the placement and the plans run on another. The forward thread only enqueues the same-GPU copies, the staging-to-weight copies and the map update for one layer between two forwards, and polls.
• Lockstep by agreement, not by a pass count: once per forward pass during a rebalance, every rank contributes a flag to a one-element all-reduce on the world group's CPU (gloo) process group, first "my placement is ready", then "my reads for layer L are done". A layer is installed on every rank at the first boundary where the all-reduce says all ranks have it, so the maps for a layer switch everywhere at the same pass and no rank can route to a slot a peer has not filled. A peer's reads of layer L's live weights are therefore always complete before anyone installs L. The worker thread waits on the previous install's event before its reads so the staging buffers are free.
• expert_location.py: ExpertLocationMetadata.compute_placement_by_eplb runs the EPLB algorithm and, with static dispatch, the per-rank dispatch map on whatever device the counts are on, returning an EplbPlacement; init_by_placement and an optional precomputed dispatch map in _init_raw; init_by_eplb is the composition of the two. compute_logical_to_rank_dispatch_physical_map takes an optional device.
• eplb_manager.py: _rebalance_async, taken by rebalance() when --eplb-async-rebalance is set and this is not an elastic-EP post-scale rebalance; on_model_loaded registers the expert tensors and the staging buffers with the engine and exchanges session ids and addresses once, after load_model, before the KV cache is sized (model_runner.py). reset_generator drains in-flight work.
• mooncake_transfer_engine.py: the shared engine is initialized when --eplb-async-rebalance is set, so the flag also works outside PD disaggregation.
• Arg --eplb-async-rebalance (default off), validated in handle_eplb_and_dispatch (needs --enable-eplb, rejects elastic EP, warns that layers-per-chunk is ignored); docs row. The synchronous path and expert_location_updater.py are untouched.
• Tests, test/registered/unit/eplb/test_eplb_async_migrator.py: the planner's four cases on eight slots over two nodes; two gloo ranks on CPU drive the migrator through one rebalance with a shared-memory transport in place of the engine, rank 1's reads delayed so the agreement is exercised; weights and maps change only at each layer's install, and never before every rank is ready.

Costs: one layer of staging buffers per rank (about 0.5 GB at DEP16 for DeepSeek-V4.1-Flash); the expert tensors registered with the engine once; one gloo all-reduce of one int per forward pass while a rebalance is in flight; NIC and NVLink bandwidth shared with the steps that overlap the reads; the worker threads share the interpreter lock during the placement.

Accuracy Tests

Not yet run on a GPU. The planned check is GSM8K under rebalances on DEP2 as in #44847, async against the synchronous path.

Speed Tests and Profiling

Not yet run on a GPU with this transport. Planned: DEP16 decode with EPLB on, async against the default single chunk, step times from the decode log around each rebalance, and the per-install cost from a trace. The previous NCCL side-stream version was run once on that shape and deadlocked as described above.

CPU-only so far: the tests above.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #39867800239
Latest PR Test (Extra): ❌ Run #39867799960
Latest PR Test (AMD ROCm 10): ❌ Run #39867799966

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit 776f60a.

Approve preview