[MengYu10151]

Motivation

DeepEP v2's ElasticBuffer is an NCCL-symmetric-memory all-to-all engine with a
fixed per-rank capacity. Fixed capacity means static communication shapes, which
makes the MoE decode path CUDA-graph capturable under any topology — including
multi-node, where legacy low_latency is not available as a graphable path.
This PR adds it as a standalone backend deepep_v2 alongside deepep; only
expert-parallel dispatch/combine is replaced. Validated on Hopper (H20, H200)
and Blackwell (B200 sm_100, B300 sm_103).

Modifications

• layers/moe/token_dispatcher/deepep_v2.py (new): the dispatcher.
  --deepep-v2-mode direct|hybrid selects the ElasticBuffer topology — not
  the legacy normal/low_latency split. Layout is chosen per batch by inference
  phase, not by comm mode: decode uses the expanded layout + masked GEMM (static
  shapes, graph-capturable, available under hybrid too), extend/prefill uses
  the contiguous layout. Single-shot dispatch()/combine() with a strict handle
  lifecycle; the guards run before the DeepEP import check so the module imports
  cleanly without DeepEP installed.
• kernels/ops/moe/ep_moe_kernels.py: masked-slab repack kernels. A fixed
  per-expert worker pool keeps the grid static (graph-safe); top-k weights are
  folded in-kernel; the scale copy takes the source pack-dim stride, which is
  what makes Blackwell's packed-UE8M0 (column-major) layout correct alongside
  Hopper's row-major fp32.
• layers/moe/moe_runner/deep_gemm.py: the FP8 adapter (masked decode /
  contiguous prefill, consuming DeepEP's 128-aligned psum directly as group
  offsets). deep_gemm is the only supported runner: deepep_v2 dispatches FP8
  activations plus scales, and no other a2a backend targets a different runner
  family either. --moe-runner-backend is checked at startup, but the runner is
  chosen per layer by the quant method and several hard-select Triton regardless,
  so MoeRunner rejects the mismatch too — otherwise such a model would reach
  the permute registry and die on a bare assert after the weights are loaded.
• layers/moe/ep_moe/layer.py: get_moe_impl_class resolves deepep_v2 to
  DeepEPMoE like every other a2a backend — this is what makes the UE8M0
  load-time weight requant apply on Blackwell.
• server_args.py, arg_groups/overrides.py: the new flags; ep_size and
  disable_shared_experts_fusion are declared through the a2a passes, and the
  runner is validated against resolved_view(...) so a model declaration that
  restores an unsupported runner fails fast. Boot-time capacity check for the
  chunked-prefill budget. TBO/SBO and enforced shared-expert fusion are rejected
  at startup.
• state_capturer/routed_experts.py: classify deepep_v2 as a scattered-a2a
  backend; otherwise dp_rank > 0 reads unwritten buffer rows under
  --enable-return-routed-experts.
• No CI/Docker changes: sgl-deep-ep==0.1.1 already provides
  ElasticBuffer, so CI can exercise this backend as-is; it stays inert for
  every other configuration, fails fast when DeepEP is missing, and its unit
  tests need only Triton.

Usage

DeepEP v2 needs NCCL symmetric memory, which SGLang turns off by default
(NCCL_CUMEM_ENABLE is set to int(enable_symm_mem) when the variable is
unset), so an otherwise correct launch fails with nccl.cu:104: Communicator does not support symmetric memory!. Pass NCCL_CUMEM_ENABLE=1 (used for every
benchmark here) or --enable-symm-mem, which also works but additionally
enables NVLS and a 4 GB symmetric-memory prealloc.

SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK is the fixed per-rank buffer
capacity. It must cover the per-rank chunk (--chunked-prefill-size / dp_size
under DP attention), and that chunk must be a multiple of page_size — with
page_size=256, dp_size=8 the smallest usable pair is capacity 256 with
--chunked-prefill-size 2048. Smaller capacity is better for decode; prefill
servers size it to their chunk budget and add --disable-cuda-graph.

No extra install step: the pinned sgl-deep-ep==0.1.1 ships ElasticBuffer
(deep_ep 2.1.0), so this backend runs on a stock environment. Re-validated on
H20 x8 against that wheel with no source build and no code change -- direct
GSM8K 0.970 / MMLU 0.890, hybrid 0.970 / 0.886, three-question gate passing in
both modes. Earlier development used a DeepEP source build at d4f41e4; the
packaged API is a superset of what the dispatcher calls.

Multi-node: add --deepep-v2-mode hybrid --enable-dp-lm-head --nnodes 2 --node-rank {0,1}. Blackwell with a BF16 wo_a checkpoint:
SGLANG_OPT_FP8_WO_A_GEMM=0.

Accuracy Tests

DeepSeek-V4-Flash-FP8, --tp 8 --dp 8 --ep 8 --enable-dp-attention,
deep_gemm, kv fp8.

	deepep_v2 direct	deepep_v2 hybrid	deepep
H20 x8 — MMLU (500)	0.888	0.890	0.888
H20 x8 — GSM8K (100)	0.980	0.980	0.970
B200 x8 — MMLU (500)	0.886	0.884	—
B200 x8 — GSM8K (100)	0.950	0.970	—

MMLU at 500 examples has a ~1.4pp standard error at this accuracy, which is
wider than every gap in the table; repeat runs of one configuration spanned
0.884-0.896. These numbers show the absence of a systematic shift against the
same-node deepep baseline, not a ranking.

The three-question gate passes on every configuration above. Unit tests (CI, no DeepEP required): masked-slab repack including a
production packed-UE8M0 scale and CUDA-graph capture/replay, handle-lifecycle
guards, server-args resolution, capturer classification, plus a DP>1
routed-experts readback e2e. The three-question gate also passes on the
multi-node configuration below. Earlier validation on the same dispatch path:
needle retrieval across chunk boundaries and 8-way concurrent requests 24/24;
prefill/decode disaggregation functional.

Speed Tests

	deepep_v2	deepep	delta
H20 x8 decode	4,339 tok/s	4,379	-0.9%
H20 x8 prefill	22,353 in_tok/s	21,651	+3.2%
B200 x8 decode	8,429 tok/s	8,675	-2.8%
B200 x8 prefill	32,447 in_tok/s	31,533	+2.9%

Each backend pair was measured back-to-back on the same node, on
torch 2.13 / sglang-kernel 0.4.6.post1, with the prefill cells interleaved so
drift cannot bias one side. H20 prefill averages 4 samples per cell (sd 95 and
72), B200 prefill 2 per cell (sd 132 and 57); decode is one cell per side.
Decode reproduces well across nodes -- deepep_v2 measured 8,385 / 8,403 / 8,429
tok/s on three different B200 nodes -- but the deepep side moved more, so the
decode delta has landed between -0.9% and -2.8% across runs; read it as "a low
single-digit percentage behind" rather than a precise figure. Prefill has been
consistently ahead by roughly 3%.

Multi-node (2x H20x8, TP16/EP16/DP16, hybrid, same protocol): 4,610 tok/s
decode with the masked-GEMM + CUDA-graph path over RDMA — a graphable cross-node
decode path legacy low_latency does not provide — and 20,688 in_tok/s
prefill (two samples: 20,322 / 21,054). There is no same-fabric deepep baseline
for this pair: DeepEP's NVSHMEM data plane does not complete across nodes on
this RoCE setup, reproduced with DeepEP's own test_internode.py outside
SGLang.

Reproduction (5 rounds per boot, mean of the last 3; prefill flushes the radix
cache before each round):

# common: NCCL_CUMEM_ENABLE=1, SGLANG_DEEPEP_V2_NUM_SMS=20, capacity 1024
#   deepep_v2: --moe-a2a-backend deepep_v2 --deepep-v2-mode direct
#   deepep   : --moe-a2a-backend deepep --deepep-mode low_latency|normal --deepep-dispatcher-output-dtype fp8
#              (with SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=1024)
NCCL_CUMEM_ENABLE=1 SGLANG_DEEPEP_V2_NUM_SMS=20 \
SGLANG_DEEPEP_V2_NUM_MAX_DISPATCH_TOKENS_PER_RANK=1024 \
python3 -m sglang.launch_server --model-path $MODEL --trust-remote-code \
  --tp-size 8 --dp-size 8 --ep-size 8 --enable-dp-attention --context-length 4096 \
  --chunked-prefill-size 8192 --moe-runner-backend deep_gemm --kv-cache-dtype fp8_e4m3 \
  --host 127.0.0.1 --port 30000 $BACKEND_FLAGS \
  --mem-fraction-static 0.70 --cuda-graph-max-bs 128                    # decode
# --mem-fraction-static 0.80 --disable-cuda-graph --disable-piecewise-cuda-graph  # prefill

# decode
python3 -m sglang.bench_serving --backend sglang --port 30000 --dataset-name random \
  --num-prompts 128 --random-input-len 1 --random-output-len 1024 --random-range-ratio 1 --max-concurrency 128
# prefill (curl -s localhost:30000/flush_cache before each round)
python3 -m sglang.bench_serving --backend sglang --port 30000 --dataset-name random \
  --num-prompts 64 --random-input-len 1024 --random-output-len 1 --random-range-ratio 1 --max-concurrency 16

Checklist

• Format code with pre-commit.
• Add unit tests.
• Update documentation (separate follow-up PR).
• Provide accuracy and speed benchmarks (above).
• Code style.

Notes for reviewers

• Prerequisites: none beyond the current pins -- sgl-deep-ep==0.1.1 provides
  ElasticBuffer. NCCL symmetric memory must be on (see Usage).
• Follow-ups: TBO/SBO comm-compute overlap (a PoC showed the naive two-phase
  split cannot overlap because ElasticBuffer.dispatch is host-blocking; it
  needs a return_recv_hook-style split in DeepEP), adapters for other MoE
  runners, documentation.

Co-authored-by: jianxiongzhang 40026647+jianxiongzhang@users.noreply.github.com

---

CI States

Latest PR Test (Base): ❌ Run #32259471892
Latest PR Test (Extra): ❌ Run #32259470116

[gemini-code-assist]

Code Review

This pull request introduces support for the DeepEP v2 (ElasticBuffer) MoE A2A backend (epv2), integrating it with both the deep_gemm and triton runners. It adds the EpV2Dispatcher, implements Triton kernels for repacking expanded expert-packed buffers into masked-GEMM slabs to optimize DeepGEMM execution, and updates model configurations, server arguments, and CI scripts accordingly. The feedback highlights several critical and high-severity issues: a missing boundary mask in the _fwd_kernel_ep_scatter_psum_init Triton kernel that can cause out-of-bounds writes, the potential propagation of NaN values due to using torch.empty instead of torch.zeros for gather_out, a robustness improvement for retrieving LOCAL_RANK during distributed initialization, and a missing import check before calling the FP8 quantization kernel in the EPv2 dispatcher.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The Triton kernel _fwd_kernel_ep_scatter_psum_init writes to m_indices without any boundary mask. Since different experts have arbitrary token counts that are not multiples of BLOCK_E, and different program instances run concurrently, this unmasked store will write out-of-bounds of the current expert's region and corrupt the indices of subsequent experts. A mask idx < cur_end must be applied to prevent non-deterministic corruption and out-of-bounds writes.

Suggested change

	off_expert = tl.arange(0, BLOCK_E)
	for start_m in tl.range(0, cur_token_num, BLOCK_E, num_stages=4):
	tl.store(m_indices + cur_start + start_m + off_expert, cur_expert)
	off_expert = tl.arange(0, BLOCK_E)
	for start_m in tl.range(0, cur_token_num, BLOCK_E, num_stages=4):
	idx = cur_start + start_m + off_expert
	tl.store(m_indices + idx, cur_expert, mask=idx < cur_end)

[gemini-code-assist]

Using torch.empty to allocate gather_out can leave unrouted or padding tokens with uninitialized garbage/NaN values. Since gather_out is the final output of the MoE layer and is added directly to the residual connection, any garbage values will propagate through the model and potentially cause NaNs. Allocating with torch.zeros is a safer and necessary defensive practice here.

Suggested change

	gather_out = torch.empty(
	running_state["hidden_states_shape"],
	device=running_state["hidden_states_device"],
	dtype=torch.bfloat16,
	)
	gather_out = torch.zeros(
	running_state["hidden_states_shape"],
	device=running_state["hidden_states_device"],
	dtype=torch.bfloat16,
	)

[gemini-code-assist]

When local_rank is -1 and the init method is not "env://", the code currently falls back to using rank directly. However, launchers like torchrun still set the LOCAL_RANK environment variable even when using other initialization methods (e.g., TCP or shared file). Checking LOCAL_RANK first regardless of the initialization method is more robust and standard.

Suggested change

	init_local_rank = local_rank
	if init_local_rank == -1:
	if distributed_init_method == "env://":
	init_local_rank = int(os.environ.get("LOCAL_RANK", "0"))
	else:
	init_local_rank = rank
	init_local_rank = local_rank
	if init_local_rank == -1:
	env_local_rank = os.environ.get("LOCAL_RANK")
	if env_local_rank is not None:
	init_local_rank = int(env_local_rank)
	else:
	init_local_rank = rank

[gemini-code-assist]

In the use_masked branch, sglang_per_token_group_quant_fp8 is called directly without calling _ensure_fp8_quant_available(). If the quantization kernel is not available (e.g., due to import errors), this will raise a confusing TypeError: 'NoneType' object is not callable instead of a clean ImportError. Calling _ensure_fp8_quant_available() at the start of the FP8 block ensures a clean fail-fast behavior.

Suggested change

	if self._uses_fp8_dispatch_output():
	if use_masked:
	if self._uses_fp8_dispatch_output():
	_ensure_fp8_quant_available()
	if use_masked:

[laixinn]

“Communicator does not support symmetric memory!” is raised with EP16 when strictly following your launch configuration. The issue is resolved by adding --enable-symm-mem.
We may be using different testbeds, but I believe DeepEP v2 should enable symmetric memory automatically rather than require users to specify this flag.

[MengYu10151]

“Communicator does not support symmetric memory!” is raised with EP16 when strictly following your launch configuration. The issue is resolved by adding --enable-symm-mem. We may be using different testbeds, but I believe DeepEP v2 should enable symmetric memory automatically rather than require users to specify this flag.

Thanks for the comments! the launch cmd was incomplete: SGLang sets NCCL_CUMEM_ENABLE to int(enable_symm_mem) when the variable is not already in the environment, so symmetric memory is off by default and DeepEP v2 fails at ElasticBuffer creation.

I've updated the reproduction commands in the description to pass NCCL_CUMEM_ENABLE=1 explicitly, which is what all the benchmarks in this PR used. Starting with --enable-symm-mem as you did works as well.