moe: add DeepEP V2 ElasticBuffer support to MoE token dispatcher#24443
moe: add DeepEP V2 ElasticBuffer support to MoE token dispatcher#24443dmvevents wants to merge 4 commits into
Conversation
DeepEP V2 (deepseek-ai/DeepEP#605, merged 2026-04-29) introduces `ElasticBuffer` alongside the legacy `Buffer`. Both classes are exported from `deep_ep.__init__`, so SGLang's existing V1 code path continues to work unchanged on a V2 install. This patch adds a second import probe (`have_deepep_v2`) next to the existing `use_deepep` flag, then lets users opt in to the V2 `ElasticBuffer` ctor behind `SGLANG_DEEPEP_USE_V2=1`. The V1 path is byte-identical when the env var is unset, so this is a backwards- compatible addition. V2 collapses the V1 NVL/RDMA byte-pool ctor into a single MoE-shape ctor and derives the internal buffer size from `num_max_tokens_per_rank`, `hidden`, and `num_topk`. `num_allocated_qps=0` asks V2 to auto-size and auto-cap the Queue-Pair budget on AWS EFA (the 128-slot GIN ring ceiling; see `_is_efa_fabric` in `deep_ep/buffers/elastic.py`). Why: - Removes the need for users to monkeypatch `deep_ep.Buffer` to reach V2 semantics, and mirrors the probe shape already used in NVIDIA/Megatron-LM's `fused_a2a.py` (`HAVE_DEEP_EP_V2`). Opt-in env-var gate keeps the change infra-bump-shaped, not a default switch flip. - vLLM is landing a parallel native-V2 path in vllm-project/vllm#41183 (16 commits, active review). SGLang users running mixed-framework MoE stacks benefit from having V2 reachable in SGLang without a runtime monkeypatch. V2 parity rules baked in: - `num_max_tokens_per_rank` sourced from `envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` when the caller did not pin it (matches the V1 low-latency path). - `num_topk=0` at ctor time: DeepEP V2 accepts this and derives a conservative group-size-based hint, so we don't need to know the router topk at buffer construction. - `clean_buffer()` skips the legacy `clean_low_latency_buffer` call when the underlying buffer is an ElasticBuffer (no equivalent on V2, per `deep_ep/buffers/elastic.py`). Related: - DeepEP V2 PR: deepseek-ai/DeepEP#605 (merged 2026-04-29) - vLLM DeepEP V2: vllm-project/vllm#41183 (open) - Megatron-LM DeepEP V2: NVIDIA/Megatron-LM#4632 (open)
Adds a lightweight guard for the `have_deepep_v2` / `use_deepep` probe pattern introduced in the previous commit. Exercises the three reachable states: - V1 `Buffer` only installed (legacy path) - V2 `ElasticBuffer` + `Buffer` both exported (typical V2 install) - `deep_ep` missing entirely The test uses a synthetic `deep_ep` module (via sys.modules monkey patch) so it runs on any CPU host without CUDA, NCCL, or EFA. The test is a `skipTest`-tolerant guard: when the SGLang module stack cannot be imported in the CI environment (for any upstream reason), it skips rather than fails, so this does not introduce a new dependency on a full SGLang install in the unit-test tier.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for DeepEP V2 by adding a probe for ElasticBuffer and an opt-in path gated by the SGLANG_DEEPEP_USE_V2 environment variable. It includes the implementation of _build_v2_buffer and a new unit test to verify the probe logic. Feedback includes suggestions to explicitly define ElasticBuffer as None on import failure for better static analysis, centralize the new environment variable within the Envs class, and remove an unused parameter in the _build_v2_buffer method.
| try: | ||
| from deep_ep import ElasticBuffer | ||
|
|
||
| have_deepep_v2 = True | ||
| except ImportError: | ||
| have_deepep_v2 = False |
There was a problem hiding this comment.
To improve code robustness and assist static analysis tools, it is recommended to explicitly define ElasticBuffer as None in the except block. This ensures the symbol is always present in the module namespace, even if the import fails.
| try: | |
| from deep_ep import ElasticBuffer | |
| have_deepep_v2 = True | |
| except ImportError: | |
| have_deepep_v2 = False | |
| try: | |
| from deep_ep import ElasticBuffer | |
| have_deepep_v2 = True | |
| except ImportError: | |
| ElasticBuffer = None | |
| have_deepep_v2 = False |
| if have_deepep_v2 and get_bool_env_var( | ||
| "SGLANG_DEEPEP_USE_V2", default="false" |
There was a problem hiding this comment.
The environment variable SGLANG_DEEPEP_USE_V2 should be added to the Envs class in python/sglang/srt/environ.py to maintain consistency with other DeepEP configurations and leverage the centralized environment management system. Using envs.SGLANG_DEEPEP_USE_V2.get() is preferred over direct calls to get_bool_env_var.
| cls, | ||
| group: dist.ProcessGroup, | ||
| hidden_size: int, | ||
| param_bytes: int, |
There was a problem hiding this comment.
The parameter param_bytes is passed to _build_v2_buffer but is not utilized within the method body. It should be removed from the function signature to improve code clarity and maintainability.
* Add minimal SGLang overlay for Wave 29 smoke Scaffold-only Dockerfile lacked pip install; would produce base-only image. Adding minimal overlay: - pip install sglang[all]==0.5.6.post2 - torch/nccl/nvshmem pin guards matching v0.2.5 base Wave 29 will materialize full adapter (sgl-project/sglang#24443) + JIT cache. * Add CodeBuild buildspec + pins.env + minimal overlay Wave 29 prep: - Minimal SGLang 0.5.6.post2 overlay (full adapter TBD) - CodeBuild buildspec + pins.env matching sibling pattern - NCCL/NVSHMEM pin guards for v0.2.5 base Unblocks Plan A Wave 29 cross-node smoke. --------- Co-authored-by: Anton Alexander <dmvevents@users.noreply.github.com>
The pre-commit lint job flagged two files that exceeded black's default 88-character line limit on multi-line if/assignment continuations. Collapsed them back onto single lines to match what black reformats them to: - python/sglang/srt/layers/moe/token_dispatcher/deepep.py - test/srt/test_deepep_v2_probe.py No logic changes. Addresses lint check failure on PR sgl-project#24443. Signed-off-by: Anton Alexander <dmvevents@users.noreply.github.com>
|
Triage of the 9 red checks (from check-run logs on head
What actually needed a code change: 1 of 9 (lint). The other 8 are all symptoms of the PR being in Draft state — sglang's CI policy is to short-circuit every downstream check until the PR is marked ready. Keeping this in Draft until we get reviewer signal on the V2 opt-in design, but the lint fix is landed so the CI will be fully green the moment this flips to Ready. Diff for the lint fix is purely formatting (-9 / +3 lines, 2 files, no logic change):
|
|
Agree on both. Pushing:
Single follow-up commit on the same branch. |
- deepep.py:68 — set ElasticBuffer = None in the except ImportError branch so static-analysis tools always see the symbol in module namespace (matches the have_deepep_v2 = False fallback shape). - deepep.py:279 — drop unused param_bytes from _build_v2_buffer. DeepEP V2's ElasticBuffer ctor derives buffer size from num_max_tokens_per_rank * hidden * num_topk; the legacy V1 byte-pool sizing is not used. The V1 path in get_buffer keeps param_bytes unchanged.
|
Pushed
CI should pick it up automatically. |
| packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = ( | ||
| buffer.low_latency_dispatch( | ||
| hidden_states, | ||
| topk_ids, | ||
| self.num_max_dispatch_tokens_per_rank, | ||
| self.num_experts, | ||
| use_fp8=use_fp8, | ||
| **(dict(use_nvfp4=True) if use_nvfp4 else dict()), | ||
| **( | ||
| dict(x_global_scale=input_global_scale) | ||
| if input_global_scale is not None | ||
| else dict() | ||
| ), | ||
| async_finish=not self.return_recv_hook, | ||
| return_recv_hook=self.return_recv_hook, | ||
| **fp8_deepgemm_scale_opts, |
There was a problem hiding this comment.
ElasticBuffer does not have a low_latency_dispatch method.
I don't think that simply changing the construction of a buffer will allow direct use of DeepEP V2.
V2 allows a hybrid mode, which can uniformly use dispatch method to handle normal and low-latency calls.
Are there any other updates?
There was a problem hiding this comment.
Correct — V2 only exposes dispatch(); low_latency_dispatch() does not exist on ElasticBuffer. This PR is the probe + opt-in ctor only — the dispatch call sites at line 565 and line 746 are still V1 and unchanged. The PR body explicitly defers full V2 dispatch/combine wiring as a follow-up ("not yet end-to-end validated in this PR"), and the PR is in Draft for that reason. Full V2 dispatch path (single unified dispatch() covering both normal and low-latency, new 5-tuple return, EPHandle lifetime for low-latency cleanup) lands in a follow-up once this probe is in. Thanks for catching it.
3 participants
Summary
Adds an opt-in DeepEP V2
ElasticBufferpath next to the existing V1Buffercode path inpython/sglang/srt/layers/moe/token_dispatcher/deepep.py. Whendeep_ep.ElasticBufferis importable ANDSGLANG_DEEPEP_USE_V2=1,DeepEPBuffer.get_deepep_buffer()constructs anElasticBufferinstead of a legacyBuffer. Otherwise the V1 path runs byte-identical. Mirrors theHAVE_DEEP_EP_V2probe pattern already used in NVIDIA/Megatron-LM'sfused_a2a.py.This is a probe-plumbing patch, not a default switch flip. It makes the V2 ctor reachable in SGLang without runtime monkeypatches, so downstream forks and users can adopt V2 semantics on their own timelines.
Motivation
DeepEP V2 (deepseek-ai/DeepEP#605) merged on 2026-04-29 and introduces
ElasticBufferalongside the legacyBuffer. Key V2 changes:ElasticBuffer(group, num_max_tokens_per_rank, hidden, num_topk, ...)instead of V1's NVL/RDMA byte-pool ctor.topk_idx, soget_dispatch_layout()is no longer required at the call site.(recv_x, recv_topk_idx, recv_topk_weights, EPHandle, event)instead of the V1 6-tuple;num_recv_tokens_per_expert_listnow lives onEPHandle.num_allocated_qps=0asks V2 to auto-size and auto-cap the Queue-Pair budget. On AWS EFA this clamps to the safe 128-slot GIN ring ceiling (see_is_efa_fabricin deep_ep/buffers/elastic.py).Critically,
deep_ep.__init__on V2 exports bothBuffer(frombuffers/legacy.py) andElasticBuffer(frombuffers/elastic.py). So SGLang's existingfrom deep_ep import Buffer, Configsurface continues to work unchanged on a V2 install. The gap this PR closes: there is currently no way for SGLang users to reachElasticBufferfrom the MoE dispatcher without monkeypatchingDeepEPBuffer.get_deepep_buffer()at runtime.Parallel upstream work on V2 MoE dispatch:
try: from deep_ep import ElasticBuffer; HAVE_DEEP_EP_V2 = True; except ImportError: ...probe shape. This PR adopts the same shape for SGLang so the three frameworks are consistent.Design choice: opt-in env var, not default switch flip
A default-on V2 path would:
ElasticBuffer— although this PR's probe makes that safe at import time, a misconfigured install would silently fall back and mask a real error.So this PR ships only the reachability plumbing (probe + opt-in ctor). Users who want V2 semantics set
SGLANG_DEEPEP_USE_V2=1at launch. Maintainers can later flip the default and replace the V1 Buffer path once V2 has been validated end-to-end on SGLang's CI.What changed
python/sglang/srt/layers/moe/token_dispatcher/deepep.pyhave_deepep_v2probe,_build_v2_bufferclassmethod, V2 branch inget_deepep_buffer,clean_bufferV2-safety guardtest/srt/test_deepep_v2_probe.pyV1 fall-through:
use_deepep=True, have_deepep_v2=False(pre-V2 DeepEP install) → legacyBufferpath runs byte-identical to the pre-patch state.use_deepep=True, have_deepep_v2=True, SGLANG_DEEPEP_USE_V2 unset→ legacyBufferpath, byte-identical to pre-patch. V2 is latent.use_deepep=True, have_deepep_v2=True, SGLANG_DEEPEP_USE_V2=1→ new V2ElasticBufferpath.use_deepep=False(no deep_ep installed) → unchanged, module still imports.V2 parity rules:
num_max_tokens_per_ranksourced fromenvs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANKwhen the caller does not pin it (matches the V1 low-latency path).num_topk=0at ctor time — V2 accepts this and derives a conservative group-size-based hint. SGLang doesn't know the router topk at buffer construction time.num_allocated_qps=0— V2 auto-sizes with an EFA safety cap (see_is_efa_fabricindeep_ep/buffers/elastic.py).clean_buffer()skipsclean_low_latency_bufferwhen the underlying object is anElasticBuffer(V2 has no equivalent; low-latency cleanup flows throughEPHandlelifetime in V2).Evidence
deep_ep/buffers/legacy.pyanddeep_ep/buffers/elastic.py.antonai-work/vllm-deepep-v2-efavalidates the shared DeepEP V2 + AWS EFA stack end-to-end (identical Dockerfile base, sameElasticBufferctor shape). The training-side equivalent isantonai-work/nemo-rl-deepep-v2-efa. Neither targets SGLang specifically — the reproducible substrate is shared; the SGLang overlay in this PR is the consumer-side plumbing.sglang-deepep-v2-efareproduction repo can be added if reviewers request it. At the probe-patch level the existing siblings already cover the V2 + EFA substrate.The SGLang V2 path itself is not yet end-to-end validated in this PR. The unit test exercises the probe plumbing; no claim is made about dispatch/combine correctness under
SGLANG_DEEPEP_USE_V2=1. That validation is a follow-up once this probe lands and the door is open.Backwards compatibility
DeepEPBuffer,DeepEPConfig,_DeepEPDispatcherImplNormal,_DeepEPDispatcherImplLowLatency,DeepEPDispatcherare unchanged.use_deepepsemantics unchanged.SGLANG_DEEPEP_USE_V2is opt-in, default false. V1 install (deep_ep < 2.0) sees only the import probe and the probe is a no-op.try/except ImportErrorat module load (microseconds). When V2 isn't installed, all V2-specific code is dead.clean_buffer()tweak is a pure safety guard — it only kicks in when the V2 path is active (noclean_low_latency_buffermethod to call).Tests
test/srt/test_deepep_v2_probe.pyexercises the three probe states via a stubdeep_epmodule. Runs on CPU-only hosts.skipTest-tolerant: if the SGLang module stack cannot import in a minimal test environment, the test skips rather than fails, so it doesn't introduce a new dependency on a full SGLang install in the unit-test tier.test/manual/ep/andtest/registered/ep/run unchanged — they use the V1 code path (default) and are unaffected by this patch.Related
antonai-work/vllm-deepep-v2-efa(inference validated)antonai-work/nemo-rl-deepep-v2-efa(training validated)/cc @zhyncs @merrymercy @Ying1123