From 65b14881c5a6d308f00958a3095dea172e7cf47b Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 9 Jul 2026 02:08:35 -0700 Subject: [PATCH] [refactor] Move the EP dispatcher and fusion-workspace manager state onto ctx.resources (#30489) --- python/sglang/srt/elastic_ep/elastic_ep.py | 2 +- .../srt/layers/flashinfer_comm_fusion.py | 36 +++++++---- .../srt/layers/moe/token_dispatcher/deepep.py | 59 ++++++++++++------ .../layers/moe/token_dispatcher/mooncake.py | 44 +++++++++---- .../srt/layers/moe/token_dispatcher/nixl.py | 61 ++++++++++++------- .../layers/test_flashinfer_comm_fusion.py | 13 +++- test/registered/unit/test_runtime_context.py | 37 +++++++++++ 7 files changed, 182 insertions(+), 70 deletions(-) diff --git a/python/sglang/srt/elastic_ep/elastic_ep.py b/python/sglang/srt/elastic_ep/elastic_ep.py index 16a854618..2490d49ff 100644 --- a/python/sglang/srt/elastic_ep/elastic_ep.py +++ b/python/sglang/srt/elastic_ep/elastic_ep.py @@ -141,7 +141,7 @@ def _maybe_create_message_queue(group) -> None: def _refresh_ep_members() -> None: from sglang.srt.layers.moe.token_dispatcher.mooncake import EPBuffer - EPBuffer._buffer.update_ep_member() + EPBuffer.get_existing_buffer().update_ep_member() def try_recover_ranks(global_ranks: List[int]) -> bool: diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 9392521d4..06abf4ca1 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -580,14 +580,22 @@ class FlashInferWorkspaceManager: self._logged_init = False -_attn_tp_workspace_manager = FlashInferWorkspaceManager() -_moe_tp_workspace_manager = FlashInferWorkspaceManager() - - def _get_workspace_manager(use_attn_tp_group: bool) -> FlashInferWorkspaceManager: - return ( - _attn_tp_workspace_manager if use_attn_tp_group else _moe_tp_workspace_manager + """The per-group fusion workspace manager; the instances live on + ``ctx.resources`` (one per comm group, created lazily).""" + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + name = ( + "flashinfer_fusion_attn_tp_workspace" + if use_attn_tp_group + else "flashinfer_fusion_moe_tp_workspace" ) + manager = buffers.get(name) + if manager is None: + manager = FlashInferWorkspaceManager() + buffers[name] = manager + return manager def _sync_allreduce_unavailable_across_tp(): @@ -853,11 +861,13 @@ def pre_initialize_workspaces( def cleanup_flashinfer_workspace(): - global _attn_tp_workspace_manager, _moe_tp_workspace_manager - if _attn_tp_workspace_manager is not None: - _attn_tp_workspace_manager.cleanup() - if ( - _moe_tp_workspace_manager is not None - and _moe_tp_workspace_manager is not _attn_tp_workspace_manager + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + for name in ( + "flashinfer_fusion_attn_tp_workspace", + "flashinfer_fusion_moe_tp_workspace", ): - _moe_tp_workspace_manager.cleanup() + manager = buffers.get(name) + if manager is not None: + manager.cleanup() diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 2d7192e2e..f0fdfe270 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -159,11 +159,27 @@ class DeepEPDispatchMode(IntEnum): class DeepEPBuffer: - _buffer = None - _dispatch_mode: Optional[DeepEPDispatchMode] = None - _hidden_size: Optional[int] = None - _num_max_dispatch_tokens_per_rank: Optional[int] = None - _num_experts: Optional[int] = None + """Managing facade for the process-wide DeepEP comm buffer; the state + itself lives on ``ctx.resources`` (one entry per process).""" + + @classmethod + def _state(cls): + from types import SimpleNamespace + + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + state = buffers.get("deepep_ep_state") + if state is None: + state = SimpleNamespace( + buffer=None, + dispatch_mode=None, + hidden_size=None, + num_max_dispatch_tokens_per_rank=None, + num_experts=None, + ) + buffers["deepep_ep_state"] = state + return state @classmethod def get_deepep_buffer( @@ -175,12 +191,13 @@ class DeepEPBuffer: num_max_dispatch_tokens_per_rank: int = -1, num_experts: int = -1, ): - if cls._buffer is not None: - return cls._buffer + state = cls._state() + if state.buffer is not None: + return state.buffer - cls._hidden_size = hidden_size - cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank - cls._num_experts = num_experts + state.hidden_size = hidden_size + state.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank + state.num_experts = num_experts num_nvl_bytes, num_rdma_bytes = 0, 0 if deepep_mode.enable_normal(): @@ -263,28 +280,30 @@ class DeepEPBuffer: if not is_cu12 and use_mnnvl_fabric: buffer_kwargs["use_fabric"] = True - cls._buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes, **buffer_kwargs) - return cls._buffer + state.buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes, **buffer_kwargs) + return state.buffer @classmethod def clean_buffer(cls): - if not cls._buffer.low_latency_mode: + state = cls._state() + if not state.buffer.low_latency_mode: return - cls._buffer.clean_low_latency_buffer( - cls._num_max_dispatch_tokens_per_rank, - cls._hidden_size, - cls._num_experts, + state.buffer.clean_low_latency_buffer( + state.num_max_dispatch_tokens_per_rank, + state.hidden_size, + state.num_experts, ) @classmethod def set_dispatch_mode_as_normal(cls): - cls._dispatch_mode = DeepEPDispatchMode.NORMAL + cls._state().dispatch_mode = DeepEPDispatchMode.NORMAL @classmethod def set_dispatch_mode_as_low_latency(cls): - if cls._dispatch_mode == DeepEPDispatchMode.NORMAL: + state = cls._state() + if state.dispatch_mode == DeepEPDispatchMode.NORMAL: cls.clean_buffer() - cls._dispatch_mode = DeepEPDispatchMode.LOW_LATENCY + state.dispatch_mode = DeepEPDispatchMode.LOW_LATENCY @classmethod def set_dispatch_mode(cls, mode: DeepEPMode): diff --git a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py index f475d69d2..5fbf494e3 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py @@ -57,10 +57,31 @@ assert isinstance(MooncakeCombineInput, CombineInput) class EPBuffer: - _buffer = None - _hidden_size: Optional[int] = None - _num_max_dispatch_tokens_per_rank: Optional[int] = None - _num_experts: Optional[int] = None + """Managing facade for the process-wide Mooncake EP buffer; the state + itself lives on ``ctx.resources``.""" + + @classmethod + def _state(cls): + from types import SimpleNamespace + + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + state = buffers.get("mooncake_ep_state") + if state is None: + state = SimpleNamespace( + buffer=None, + hidden_size=None, + num_max_dispatch_tokens_per_rank=None, + num_experts=None, + ) + buffers["mooncake_ep_state"] = state + return state + + @classmethod + def get_existing_buffer(cls): + """The already-created buffer (elastic-EP membership refresh).""" + return cls._state().buffer @classmethod def get_ep_buffer( @@ -72,15 +93,16 @@ class EPBuffer: num_max_dispatch_tokens_per_rank: int = -1, num_experts: int = -1, ): - if cls._buffer is not None: - return cls._buffer + state = cls._state() + if state.buffer is not None: + return state.buffer # Lazy import Buffer to avoid creating CUDA context at module import time from mooncake.mooncake_ep_buffer import Buffer - cls._hidden_size = hidden_size - cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank - cls._num_experts = num_experts + state.hidden_size = hidden_size + state.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank + state.num_experts = num_experts num_ep_buffer_bytes = 0 if deepep_mode.enable_normal(): @@ -97,8 +119,8 @@ class EPBuffer: num_experts, ) - cls._buffer = Buffer(group, num_ep_buffer_bytes) - return cls._buffer + state.buffer = Buffer(group, num_ep_buffer_bytes) + return state.buffer class _MooncakeEPDispatcherImpl: diff --git a/python/sglang/srt/layers/moe/token_dispatcher/nixl.py b/python/sglang/srt/layers/moe/token_dispatcher/nixl.py index e1977f362..dd04f7d09 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/nixl.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/nixl.py @@ -2,7 +2,6 @@ from __future__ import annotations import logging from enum import Enum, auto -from typing import Optional import torch import torch.distributed as dist @@ -39,11 +38,27 @@ NixlEPCombineInput = DeepEPLLCombineInput class NixlEPBuffer: - _buffer = None - _hidden_size: Optional[int] = None - _num_max_dispatch_tokens_per_rank: Optional[int] = None - _num_experts: Optional[int] = None - _num_local_experts: Optional[int] = None + """Managing facade for the process-wide NIXL EP buffer; the state itself + lives on ``ctx.resources``.""" + + @classmethod + def _state(cls): + from types import SimpleNamespace + + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + state = buffers.get("nixl_ep_state") + if state is None: + state = SimpleNamespace( + buffer=None, + hidden_size=None, + num_max_dispatch_tokens_per_rank=None, + num_experts=None, + num_local_experts=None, + ) + buffers["nixl_ep_state"] = state + return state @classmethod def get_nixl_buffer( @@ -55,13 +70,14 @@ class NixlEPBuffer: num_experts: int = -1, num_local_experts: int = -1, ): - if cls._buffer is not None: - return cls._buffer + state = cls._state() + if state.buffer is not None: + return state.buffer - cls._hidden_size = hidden_size - cls._num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank - cls._num_experts = num_experts - cls._num_local_experts = num_local_experts + state.hidden_size = hidden_size + state.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank + state.num_experts = num_experts + state.num_local_experts = num_local_experts num_rdma_bytes = 0 if deepep_mode.enable_normal(): @@ -89,30 +105,31 @@ class NixlEPBuffer: logger.info( f"Using NIXL EP (world_size={world_size}, rank={rank}, " - f"num_experts={cls._num_experts}, num_experts_per_rank={cls._num_local_experts}) " + f"num_experts={state.num_experts}, num_experts_per_rank={state.num_local_experts}) " ) - cls._buffer = Buffer( + state.buffer = Buffer( rank=rank, tcp_store_group=tcp_store, ) - cls._buffer.update_memory_buffers( + state.buffer.update_memory_buffers( num_ranks=world_size, - num_experts_per_rank=cls._num_local_experts, + num_experts_per_rank=state.num_local_experts, num_rdma_bytes=num_rdma_bytes, ) all_ranks = list(range(world_size)) - cls._buffer.connect_ranks(all_ranks) + state.buffer.connect_ranks(all_ranks) - return cls._buffer + return state.buffer @classmethod def clean_buffer(cls): - cls._buffer.clean_buffer( - cls._num_max_dispatch_tokens_per_rank, - cls._hidden_size, - cls._num_experts, + state = cls._state() + state.buffer.clean_buffer( + state.num_max_dispatch_tokens_per_rank, + state.hidden_size, + state.num_experts, ) diff --git a/test/registered/unit/layers/test_flashinfer_comm_fusion.py b/test/registered/unit/layers/test_flashinfer_comm_fusion.py index 5fb1dcc7a..c2ad459fd 100644 --- a/test/registered/unit/layers/test_flashinfer_comm_fusion.py +++ b/test/registered/unit/layers/test_flashinfer_comm_fusion.py @@ -174,8 +174,12 @@ class TestFlashInferCommFusion(unittest.TestCase): fake_comm = _FakeFlashInferComm() original_comm = fusion._flashinfer_comm original_create = fusion._create_allreduce_fusion_workspace - original_manager = fusion._attn_tp_workspace_manager original_unavailable = fusion._flashinfer_allreduce_unavailable + from sglang.srt.runtime_context import get_resources + + buffers = get_resources().buffers + manager_key = "flashinfer_fusion_attn_tp_workspace" + original_manager = buffers.get(manager_key) try: fusion._flashinfer_comm = fake_comm fusion._create_allreduce_fusion_workspace = ( @@ -189,7 +193,7 @@ class TestFlashInferCommFusion(unittest.TestCase): manager = fusion.FlashInferWorkspaceManager() manager.workspace = _FakeWorkspace(backend, world_size) manager.initialized = True - fusion._attn_tp_workspace_manager = manager + buffers[manager_key] = manager if not torch.cuda.is_available(): self.skipTest("FlashInfer allreduce custom op is CUDA-only") device = torch.device("cuda") @@ -229,7 +233,10 @@ class TestFlashInferCommFusion(unittest.TestCase): finally: fusion._flashinfer_comm = original_comm fusion._create_allreduce_fusion_workspace = original_create - fusion._attn_tp_workspace_manager = original_manager + if original_manager is None: + buffers.pop(manager_key, None) + else: + buffers[manager_key] = original_manager fusion._flashinfer_allreduce_unavailable = original_unavailable diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index b903b6a6e..f847597e6 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -463,6 +463,43 @@ class TestNamedStreams(_IsolatedServerArgs): self.assertEqual(get_context().resources.streams, {}) +class TestEpBufferState(_IsolatedServerArgs): + """EP dispatcher buffer managers: state lives on ctx.resources; the + facade keeps the mode-transition and clean semantics.""" + + def test_deepep_dispatch_mode_transitions_and_reset(self): + try: + from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer + except ImportError: + self.skipTest("deep_ep not installed") + + reset_context() + cleans = [] + + class _FakeBuffer: + low_latency_mode = True + + def clean_low_latency_buffer(self, *args): + cleans.append(args) + + state = DeepEPBuffer._state() + state.buffer = _FakeBuffer() + state.hidden_size = 7168 + state.num_max_dispatch_tokens_per_rank = 128 + state.num_experts = 256 + + DeepEPBuffer.set_dispatch_mode_as_normal() + # NORMAL -> LOW_LATENCY must clean the low-latency buffer once. + DeepEPBuffer.set_dispatch_mode_as_low_latency() + self.assertEqual(cleans, [(128, 7168, 256)]) + # LOW_LATENCY -> LOW_LATENCY must not clean again. + DeepEPBuffer.set_dispatch_mode_as_low_latency() + self.assertEqual(len(cleans), 1) + + reset_context() + self.assertIsNone(DeepEPBuffer._state().buffer) + + class TestPublishLifecycle(_IsolatedServerArgs): """Publish installs the resolved server_args and seeds the capture tier."""