[refactor] Move the EP dispatcher and fusion-workspace manager state onto ctx.resources (#30489)

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
 
 

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."""