[Perf][EPLB][DSv4] Add a symm_mem weight-migration backend and fuse MegaMoE replica mapping

[abmfy]

Overview

This PR makes EPLB cheaper for DeepSeek-V4.1-Flash, running DP attention with EP MegaMoE on multi-node NVLink. It does three things:

• adds an opt-in symm_mem weight-migration backend;
• fuses EPLB's logical→physical replica mapping and expert-load recording into the DeepGEMM MegaMoE staging kernels;
• takes async EPLB's host work off the decode critical path.

Not a duplicate. No open or merged PR adds a symmetric-memory EPLB backend or fuses the MegaMoE mapping. #61176 and #61839 enable EPLB for the ROCm AITER MegaMoE backend. Related: #61767 (cpu_all_reduce_, reused here), #57549, #56716, #61643.

Claims

• EPLB raises decode throughput on DeepSeek-V4.1-Flash on GB200 NVL72 by +3.7% at DEP32 with 72 req/rank and +7.7% at DEP32 with 128 req/rank. At DEP16 the step is about 3% shorter, but the throughput gain is within run-to-run noise.
• symm_mem: senders push expert rows into the receiver's expert_buffer, which is allocated in CUDA symmetric memory. There is no staging buffer, and it needs no memory beyond expert_buffer. Cross-node peer copies run at about 800 GB/s, against about 88 GB/s for NIXL, and a 40-layer rearrangement transfers in 0.3–0.6 s at DEP32.
• MegaMoE layers no longer launch a separate EPLB mapping kernel, which saves 5–6 µs per layer. The replica choice is shared with the router through one Triton helper, so physical ids and load counts match eplb_map_to_physical_and_record exactly.
• Async EPLB no longer adds a GPU→host sync to the decode loop. It also raises on every rank when the worker fails on any rank, instead of stalling.

Validation

Hardware: GB200 NVL72, 4 GPUs per node.

Serving setup: DSpark with 5 draft tokens, MRV2, FULL CUDA graphs, async scheduling.

EPLB setup: {"num_redundant_experts": <DP>, "window_size": 1000, "step_interval": 3000, "use_async": true, "communicator": "symm_mem"}, i.e. one redundant expert per rank.

Workload: decode only, coding-agent prompts with mean ISL ≈ 85k (prefix-cached), 16k output tokens with ignore_eos, and requests pinned evenly to DP ranks. Throughput is the output rate over the window in which every request is decoding.

Off and on arms ran interleaved (ABAB) on the same nodes. These runs used a branch of nightly fe4df3d (2026-10-12) with this change set and #61767, before the rebase and without the second commit.

Config	EPLB off (tok/s)	EPLB on (tok/s)	Δ	Welch t
DEP16, 72 req/rank	223.6k ± 2.0k (n=4)	228.1k ± 4.2k (n=4)	+2.0%	2.0 (not resolved)
DEP32, 72 req/rank	419.5k ± 1.3k (n=4)	435.1k ± 5.8k (n=4)	+3.7%	5.3
DEP32, 128 req/rank	500.7k ± 5.1k (n=3)	539.4k ± 9.7k (n=3)	+7.7%	6.1

• DEP16 step time, p50 of 2000-step windows at full batch, same build: 24.34 / 24.28 ms off vs 23.61 / 23.25 ms on.
• At DEP32 with 128 req/rank, one pass per arm is excluded. A host stall unrelated to EPLB cut that pass by 13–34%.

Per-step breakdown, DEP32 with 128 req/rank:

	EPLB off	EPLB on
Step time (CUDA events, unprofiled)	38.8 ms	36.5 ms
GPU idle at the step boundary	0.27 ms	0.27 ms
Main-model MegaMoE (5 s CUDA profile)	18.1 ms	13.4 ms
All other kernels (same profile)	21.5 ms	21.2 ms

The GPU never waits on the host, so the MegaMoE saving turns into step time. MegaMoE is about 40% of the step, which caps the end-to-end gain.

Microbenchmarks (GB200):

• Fused EPLB staging, one MoE layer, hidden 5120, top-k 6, 384 logical experts. Each arm is timed as a CUDA graph replay, median of 200. Outputs, ids and load counts match main bit for bit in every case.

  Format	Tokens	Physical experts	main: torch.where + map kernel + staging	Fused	Saved
  MXFP8	432	400	9.4 µs	4.2 µs	5.1 µs
  MXFP8	768	416	11.0 µs	4.7 µs	6.3 µs
  NVFP4	432	400	49.5 µs	44.0 µs	5.5 µs
  NVFP4	768	416	49.1 µs	43.8 µs	5.3 µs

  That is about 0.2 ms per decode step across 40 MoE layers.

• Weight-migration bandwidth, 2 nodes, one MegaMoE layer of 25 experts (470 MB):

  Path	Bandwidth
  symm_mem peer copies, intra- or cross-node over multi-node NVLink	~808 GB/s (bulk), ~294 GB/s (100 per-tensor copies)
  NIXL, cross-node	~88 GB/s (UCX picks RoCE, not multi-node NVLink)
  NIXL, cross-node, UCX_TLS=cuda_copy,cuda_ipc,tcp	~0.12 GB/s (TCP)

• MegaMoE kernel time under serving load, DEP16 with 72 req/rank, layers 2–40: p50 per call drops 12.6–14.0% across three profiled runs (222→194 µs). Layer 1 is excluded because it absorbs cross-rank arrival skew.

Tests on this branch, one GB200 node with 4 GPUs. They ran before the rebase onto 64175d3; the 5 newer commits on main touch none of these files.

Test	Result
tests/kernels/moe/test_deep_gemm_mega_moe.py (incl. test_prepare_megamoe_eplb_matches_router_under_graph_replay)	110 passed
tests/kernels/moe/test_routing.py -k eplb	18 passed
tests/distributed/test_eplb_execute.py	35 passed, 14 skipped (NIXL not installed, XPU)
tests/distributed/test_eplb_execute.py -k nixl, NIXL installed	7 passed
tests/distributed/test_eplb_utils.py, test_eplb_migration_scheduler.py	6 passed, 28 passed
tests/models/test_deepseek_v41_replay_batch.py, test_deepseek_v41_decoder_replay_layers.py	13 passed, 7 passed

GSM8K on DeepSeek-V4.1-Flash, DEP4 (one node):

• Setup: 5-shot, 1319 questions, temperature 0, num_redundant_experts=16, window_size=100, step_interval=300. Every arm rearranged all 40 layers four times during the eval.
• Command: python tests/evals/gsm8k/gsm8k_eval.py --port 8000 --num-shots 5 --num-questions 1319 --max-concurrency 64

Arm	Accuracy (2 passes)	40-layer rearrangement
EPLB off	0.898, 0.901	—
symm_mem	0.897, 0.895	2.9–3.3 s
symm_mem, VLLM_GPU_SYNC_CHECK=error	0.889, 0.883	2.9–3.3 s, no sync error
nixl	0.892, 0.886	2.7–3.2 s

The EPLB-off and nixl arms ran on the commit before the final cleanup, which does not touch either path. Within one node, symm_mem and NIXL (CUDA IPC) move weights at a similar rate; symm_mem's advantage is across nodes over multi-node NVLink.

Details

symm_mem backend

• The receive buffer is one symmetric uint8 buffer per EP group, set up with a single rendezvous and carved into per-tensor views. The backend checks the NVLink/MNNVL domain before allocating.
• Senders and receivers derive each receive row from the global placement with the same helper, so add_recv is a no-op. Push backends get every rank's rows through a new recv_rows argument to set_transfer_context. This argument is optional and other backends ignore it.
• Each execute() runs two CPU barriers. GPU spin barriers can deadlock against MegaMoE's cooperative launches.
• Rejected with elastic EP. The buffer lives outside the caching allocator, so sleep mode does not free it.

Fused mapping: fused_moe/eplb_triton.py holds the replica hash, map lookup and gated load count. The MXFP8/BF16 and NVFP4 staging kernels call it, and so does the router kernel. Routed-experts capture still sees logical ids.

Async EPLB

• Per-layer maps are committed with pinned non-blocking copies, and the worker thread fences expert_buffer reuse.
• Balancedness stats are read back non-blocking and logged in order once they land.
• should_record_tensor is rewritten only when its value flips.
• The per-step "all ranks ready" check runs over cpu_all_reduce_. A gloo ring all-reduce costs several ms per step at DP16/DP32.

Also

• logical_to_physical_map is sized to R+1 slots unless elastic EP is on.
• compute_logical_maps is vectorized.
• DeepSeek V4.1 decoder-replay padding and dummy rows are marked as padding, which complements [Bugfix][MoE] Keep routing rows during DeepSeek V4.1 decoder replay #61643.

CI: the symm_mem execution tests need NVLink peers. They are added to the optional 8×H100 "Distributed DP + EP" job; the 4×L4 EPLB job skips them. The MegaMoE parity test needs SM100.

Limitations

• In mixed prefill/decode serving, per-rank prefill imbalance bounds the step. There, EPLB mostly lowers tail TTFT rather than raising throughput.
• Likely rebase overlap with [Bugfix][EPLB] Move 1-byte expert tensors as uint8 views #61786, [EPLB] Add the MoE Load Balancer policy: pluggable expert placement and replica routing #59168 and [EPLB] Add offline expert mapping support (aka "offline EPLB") #41141. Offline placements without elastic EP ([EPLB] Add offline expert mapping support (aka "offline EPLB") #41141) could exceed R+1 replicas.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--62009.org.readthedocs.build/en/62009/

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.