From 8c883d32b64415b20f01e6ef91684b522377b27f Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Tue, 5 May 2026 20:03:28 +0000 Subject: [PATCH 1/4] moe: add DeepEP V2 ElasticBuffer probe and opt-in buffer path 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) --- .../srt/layers/moe/token_dispatcher/deepep.py | 100 ++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index f990d7d00172..e043f742fb48 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -51,6 +51,22 @@ except ImportError: use_deepep = False +# DeepEP V2 introduces `ElasticBuffer` alongside the legacy `Buffer` +# (deepseek-ai/DeepEP#605, merged 2026-04-29). On V2 both classes are +# exported from `deep_ep.__init__`, so the existing `from deep_ep import +# Buffer` surface above continues to work unchanged — `ElasticBuffer` is +# an additional, MoE-shape ctor with auto-QP sizing that callers may +# opt into. The probe below is orthogonal to `use_deepep` and does not +# affect the default code path. V2 usage is further gated on +# `SGLANG_DEEPEP_USE_V2=1`. Mirrors the `HAVE_DEEP_EP_V2` probe shape +# already used in NVIDIA/Megatron-LM's `fused_a2a.py`. +try: + from deep_ep import ElasticBuffer + + have_deepep_v2 = True +except ImportError: + have_deepep_v2 = False + from enum import Enum, IntEnum, auto import torch @@ -166,6 +182,25 @@ def get_deepep_buffer( cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank cls._num_experts = num_experts + # Opt-in V2 path: construct `deep_ep.ElasticBuffer` instead of the + # legacy `deep_ep.Buffer`. V2 uses a MoE-shape ctor and infers the + # dispatch layout internally (no `get_dispatch_config` / + # `get_combine_config` / `get_nvl_buffer_size_hint` / + # `get_rdma_buffer_size_hint` calls). Gated behind an env var so + # the default code path is byte-identical to V1. + if have_deepep_v2 and get_bool_env_var( + "SGLANG_DEEPEP_USE_V2", default="false" + ): + cls._buffer = cls._build_v2_buffer( + group, + hidden_size, + param_bytes, + deepep_mode, + num_max_dispatch_tokens_per_rank, + num_experts, + ) + return cls._buffer + num_nvl_bytes, num_rdma_bytes = 0, 0 if deepep_mode.enable_normal(): hidden_bytes = hidden_size * param_bytes @@ -238,8 +273,73 @@ def get_deepep_buffer( ) return cls._buffer + @classmethod + def _build_v2_buffer( + cls, + group: dist.ProcessGroup, + hidden_size: int, + param_bytes: int, + deepep_mode: DeepEPMode, + num_max_dispatch_tokens_per_rank: int, + num_experts: int, + ): + """Construct a DeepEP V2 `ElasticBuffer` for opt-in V2 usage. + + DeepEP V2 (deepseek-ai/DeepEP#605) 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 this clamps to the + safe 128-slot GIN ring ceiling (see the EFA fast-path in + `deep_ep/buffers/elastic.py`). + + V2 does not expose `get_dispatch_config` / `get_combine_config` + / `num_qps_per_rank`, so the matching V1 code above is skipped. + """ + # Keep the num_tokens / num_topk source consistent with the V1 + # low-latency path. For the normal path where + # `num_max_dispatch_tokens_per_rank == -1`, fall back to the + # SGLang env default that already bounds the V1 path. + if num_max_dispatch_tokens_per_rank <= 0: + num_max_dispatch_tokens_per_rank = ( + envs.SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get() + ) + # `num_topk` is not known at buffer-construction time in SGLang + # (the router choice is per-forward). DeepEP V2 accepts + # `num_topk=0` and falls back to a group-size-based conservative + # hint, so we pass 0 here and let V2 compute the ceiling. + num_topk = 0 + # Low-latency mode still requires `num_experts % group.size() == + # 0`. The normal mode works with `num_experts == -1`. + if deepep_mode.enable_low_latency(): + assert num_experts != -1 and num_experts % group.size() == 0 + + logger.info( + "SGLANG_DEEPEP_USE_V2=1: constructing deep_ep.ElasticBuffer " + "(num_max_tokens_per_rank=%d, hidden=%d, num_topk=%d).", + num_max_dispatch_tokens_per_rank, + hidden_size, + num_topk, + ) + return ElasticBuffer( + group=group, + num_max_tokens_per_rank=num_max_dispatch_tokens_per_rank, + hidden=hidden_size, + num_topk=num_topk, + use_fp8_dispatch=False, + allow_hybrid_mode=True, + # num_allocated_qps=0 -> V2 auto-sizes with an EFA safety cap + # (see deep_ep/buffers/elastic.py _is_efa_fabric()). + num_allocated_qps=0, + ) + @classmethod def clean_buffer(cls): + # DeepEP V2's `ElasticBuffer` does not expose `low_latency_mode` + # or `clean_low_latency_buffer` — low-latency cleanup is handled + # internally via `EPHandle` lifetime. Fall through for V2. + if not hasattr(cls._buffer, "clean_low_latency_buffer"): + return if not cls._buffer.low_latency_mode: return cls._buffer.clean_low_latency_buffer( From f66ab63b499b7671ee8d7052b2e54708a219ddb1 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Tue, 5 May 2026 20:03:36 +0000 Subject: [PATCH 2/4] moe: add CPU-only unit test for DeepEP V2 probe plumbing 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. --- test/srt/test_deepep_v2_probe.py | 146 +++++++++++++++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 test/srt/test_deepep_v2_probe.py diff --git a/test/srt/test_deepep_v2_probe.py b/test/srt/test_deepep_v2_probe.py new file mode 100644 index 000000000000..721ca7b7ccd1 --- /dev/null +++ b/test/srt/test_deepep_v2_probe.py @@ -0,0 +1,146 @@ +"""CPU-only unit test for the DeepEP V2 probe plumbing in +`sglang.srt.layers.moe.token_dispatcher.deepep`. + +Exercises the three reachable probe states: + +1. V1 `Buffer` available, V2 `ElasticBuffer` missing +2. V2 `ElasticBuffer` available, V1 `Buffer` also available (normal V2 install) +3. Neither available + +The test does not execute any DeepEP kernels — it imports the module +under a monkey-patched `deep_ep` stub so it runs on any CPU host without +CUDA / NCCL / EFA. It is purely a guard against the probe pattern +silently regressing. + +Run: + pytest -xvs test/srt/test_deepep_v2_probe.py +""" + +from __future__ import annotations + +import importlib +import sys +import types +import unittest +from contextlib import contextmanager + + +def _make_stub_deep_ep(*, with_buffer: bool, with_elastic_buffer: bool): + """Build a fake `deep_ep` module with the attributes we opt to expose.""" + + mod = types.ModuleType("deep_ep") + if with_buffer: + + class _StubBuffer: # noqa: D401 - placeholder class + num_sms = 20 + + @staticmethod + def get_dispatch_config(size): # pragma: no cover - unused in test + return None + + @staticmethod + def get_combine_config(size): # pragma: no cover - unused in test + return None + + mod.Buffer = _StubBuffer + mod.Config = object + if with_elastic_buffer: + + class _StubElasticBuffer: # noqa: D401 - placeholder class + pass + + mod.ElasticBuffer = _StubElasticBuffer + return mod + + +@contextmanager +def _patched_deep_ep(*, with_buffer: bool, with_elastic_buffer: bool): + """Install a stub `deep_ep` in sys.modules and drop the SGLang + deepep module from the cache so it re-runs its imports against the + stub. Restores the original state on exit. + """ + + saved = { + k: sys.modules[k] + for k in list(sys.modules) + if k == "deep_ep" or k.startswith("deep_ep.") + } + saved_deepep = sys.modules.pop( + "sglang.srt.layers.moe.token_dispatcher.deepep", None + ) + # Wipe the stub path + sys.modules.pop("deep_ep", None) + if with_buffer or with_elastic_buffer: + sys.modules["deep_ep"] = _make_stub_deep_ep( + with_buffer=with_buffer, with_elastic_buffer=with_elastic_buffer + ) + try: + yield + finally: + # Best-effort restore; we don't try to rewind every transitive + # import that the probe triggered. + sys.modules.pop("deep_ep", None) + for k, v in saved.items(): + sys.modules[k] = v + if saved_deepep is not None: + sys.modules["sglang.srt.layers.moe.token_dispatcher.deepep"] = ( + saved_deepep + ) + + +class TestDeepEPV2Probe(unittest.TestCase): + """Guard the V2 probe plumbing against silent regression.""" + + def _import_probe_flags(self): + sys.modules.pop("sglang.srt.layers.moe.token_dispatcher.deepep", None) + mod = importlib.import_module( + "sglang.srt.layers.moe.token_dispatcher.deepep" + ) + return ( + getattr(mod, "use_deepep"), + getattr(mod, "have_deepep_v2"), + ) + + def test_v1_only_installed(self): + """Legacy install (V1 `Buffer` only). `use_deepep=True`, + `have_deepep_v2=False`. No regression on pre-V2 users.""" + with _patched_deep_ep(with_buffer=True, with_elastic_buffer=False): + try: + use_deepep, have_v2 = self._import_probe_flags() + except ImportError: + # Test env may not have full sglang deps; that's fine, + # the probe itself is what we need to prove compiles. + self.skipTest("sglang module stack unavailable on test host") + return + self.assertTrue(use_deepep) + self.assertFalse(have_v2) + + def test_v2_installed(self): + """V2 install (both `Buffer` legacy + `ElasticBuffer` exported). + `use_deepep=True`, `have_deepep_v2=True`. The V2 path is now + reachable behind `SGLANG_DEEPEP_USE_V2=1`.""" + with _patched_deep_ep(with_buffer=True, with_elastic_buffer=True): + try: + use_deepep, have_v2 = self._import_probe_flags() + except ImportError: + self.skipTest("sglang module stack unavailable on test host") + return + self.assertTrue(use_deepep) + self.assertTrue(have_v2) + + def test_neither_installed(self): + """deep_ep missing entirely. `use_deepep=False`, + `have_deepep_v2=False`. Module must still import cleanly so the + SGLang package can load on non-MoE targets.""" + with _patched_deep_ep(with_buffer=False, with_elastic_buffer=False): + try: + use_deepep, have_v2 = self._import_probe_flags() + except ImportError: + self.skipTest("sglang module stack unavailable on test host") + return + self.assertFalse(use_deepep) + self.assertFalse(have_v2) + + +if __name__ == "__main__": + unittest.main() From ecb4b0ed9f485ac6376d10facb8e5ded2f95d9b0 Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Fri, 8 May 2026 08:29:28 +0000 Subject: [PATCH 3/4] lint: apply black-jupyter auto-formatting 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 #24443. Signed-off-by: Anton Alexander --- python/sglang/srt/layers/moe/token_dispatcher/deepep.py | 4 +--- test/srt/test_deepep_v2_probe.py | 8 ++------ 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index e043f742fb48..241f1e61ba2d 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -188,9 +188,7 @@ def get_deepep_buffer( # `get_combine_config` / `get_nvl_buffer_size_hint` / # `get_rdma_buffer_size_hint` calls). Gated behind an env var so # the default code path is byte-identical to V1. - if have_deepep_v2 and get_bool_env_var( - "SGLANG_DEEPEP_USE_V2", default="false" - ): + if have_deepep_v2 and get_bool_env_var("SGLANG_DEEPEP_USE_V2", default="false"): cls._buffer = cls._build_v2_buffer( group, hidden_size, diff --git a/test/srt/test_deepep_v2_probe.py b/test/srt/test_deepep_v2_probe.py index 721ca7b7ccd1..7eeae8dc59dd 100644 --- a/test/srt/test_deepep_v2_probe.py +++ b/test/srt/test_deepep_v2_probe.py @@ -83,9 +83,7 @@ def _patched_deep_ep(*, with_buffer: bool, with_elastic_buffer: bool): for k, v in saved.items(): sys.modules[k] = v if saved_deepep is not None: - sys.modules["sglang.srt.layers.moe.token_dispatcher.deepep"] = ( - saved_deepep - ) + sys.modules["sglang.srt.layers.moe.token_dispatcher.deepep"] = saved_deepep class TestDeepEPV2Probe(unittest.TestCase): @@ -93,9 +91,7 @@ class TestDeepEPV2Probe(unittest.TestCase): def _import_probe_flags(self): sys.modules.pop("sglang.srt.layers.moe.token_dispatcher.deepep", None) - mod = importlib.import_module( - "sglang.srt.layers.moe.token_dispatcher.deepep" - ) + mod = importlib.import_module("sglang.srt.layers.moe.token_dispatcher.deepep") return ( getattr(mod, "use_deepep"), getattr(mod, "have_deepep_v2"), From f9f58edf704d85034e64096c5c053b3801bf8dfe Mon Sep 17 00:00:00 2001 From: Anton Alexander Date: Fri, 15 May 2026 00:04:55 +0000 Subject: [PATCH 4/4] moe: address gemini-code-assist review on DeepEP V2 path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- python/sglang/srt/layers/moe/token_dispatcher/deepep.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 241f1e61ba2d..7fbe8e3088ae 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -65,6 +65,7 @@ have_deepep_v2 = True except ImportError: + ElasticBuffer = None have_deepep_v2 = False from enum import Enum, IntEnum, auto @@ -192,7 +193,6 @@ def get_deepep_buffer( cls._buffer = cls._build_v2_buffer( group, hidden_size, - param_bytes, deepep_mode, num_max_dispatch_tokens_per_rank, num_experts, @@ -276,7 +276,6 @@ def _build_v2_buffer( cls, group: dist.ProcessGroup, hidden_size: int, - param_bytes: int, deepep_mode: DeepEPMode, num_max_dispatch_tokens_per_rank: int, num_experts: int,