From 9f5cd80a8d21899b7c1d7cd26ba7a5f83a26422e Mon Sep 17 00:00:00 2001 From: Xun Sun Date: Mon, 12 Jan 2026 17:01:24 +0800 Subject: [PATCH] Re-introduce the unit test of test_mooncake_ep_small (#16019) --- .../srt/eplb/eplb_algorithms/__init__.py | 2 +- .../layers/moe/token_dispatcher/mooncake.py | 26 +++++++------------ .../ep/test_mooncake_ep_small.py | 7 +++-- test/srt/run_suite.py | 3 +-- 4 files changed, 17 insertions(+), 21 deletions(-) rename test/{manual => srt}/ep/test_mooncake_ep_small.py (93%) diff --git a/python/sglang/srt/eplb/eplb_algorithms/__init__.py b/python/sglang/srt/eplb/eplb_algorithms/__init__.py index fc4d8f0f8..b09a14175 100644 --- a/python/sglang/srt/eplb/eplb_algorithms/__init__.py +++ b/python/sglang/srt/eplb/eplb_algorithms/__init__.py @@ -54,7 +54,7 @@ def rebalance_experts( num_groups=num_groups, num_nodes=num_nodes, num_gpus=num_physical_experts // num_local_physical_experts, - enable_hierarchical=True, + enable_hierarchical=False, active_ranks=( ElasticEPStateManager.instance().active_ranks if ElasticEPStateManager.instance() is not None diff --git a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py index c565ab855..f475d69d2 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/mooncake.py @@ -2,7 +2,11 @@ from __future__ import annotations import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, NamedTuple, Optional +from enum import Enum, auto +from typing import NamedTuple, Optional + +import torch +import torch.distributed as dist from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder @@ -18,14 +22,6 @@ from sglang.srt.layers.moe.topk import TopKOutput from sglang.srt.layers.moe.utils import DeepEPMode from sglang.srt.utils import get_int_env_var -if TYPE_CHECKING: - from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs - -from enum import Enum, auto - -import torch -import torch.distributed as dist - logger = logging.getLogger(__name__) @@ -147,8 +143,6 @@ class _MooncakeEPDispatcherImpl: self.first_execution = True self.timeout_us = 10000000 - self.active_ranks = ElasticEPStateManager.instance().active_ranks - self.handle = None def dispatch_a( @@ -215,11 +209,12 @@ class _MooncakeEPDispatcherImpl: use_fp8: bool = False, ): buffer = self._get_buffer() + active_ranks = ElasticEPStateManager.instance().active_ranks packed_recv_hidden, packed_recv_count, self.handle, event, hook = ( buffer.dispatch( hidden_states, topk_ids, - self.active_ranks, + active_ranks, self.num_max_dispatch_tokens_per_rank, self.num_experts, -1 if self.first_execution else self.timeout_us, @@ -235,14 +230,13 @@ class _MooncakeEPDispatcherImpl: hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, - overlap_args: Optional[CombineOverlapArgs] = None, ): hidden_states, event, hook = self._combine_core( hidden_states, topk_ids, topk_weights, ) - return hidden_states, event, hook, overlap_args + return hidden_states, event, hook def combine_b(self, hidden_states, event, hook): hook() if self.return_recv_hook else event.current_stream_wait() @@ -255,11 +249,12 @@ class _MooncakeEPDispatcherImpl: topk_weights: torch.Tensor, ): buffer = self._get_buffer() + active_ranks = ElasticEPStateManager.instance().active_ranks combined_hidden_states, event, hook = buffer.combine( hidden_states, topk_ids, topk_weights, - self.active_ranks, + active_ranks, -1 if self.first_execution else self.timeout_us, self.handle, async_finish=not self.return_recv_hook, @@ -368,7 +363,6 @@ class MooncakeEPDispatcher(BaseDispatcher): hidden_states=hidden_states, topk_ids=topk_ids, topk_weights=topk_weights, - overlap_args=self.overlap_args, ) self._combine_intermediate_state = inner_state diff --git a/test/manual/ep/test_mooncake_ep_small.py b/test/srt/ep/test_mooncake_ep_small.py similarity index 93% rename from test/manual/ep/test_mooncake_ep_small.py rename to test/srt/ep/test_mooncake_ep_small.py index 6e680041d..4f6d02a39 100644 --- a/test/manual/ep/test_mooncake_ep_small.py +++ b/test/srt/ep/test_mooncake_ep_small.py @@ -35,7 +35,7 @@ class TestTP(CustomTestCase): "--mooncake-ib-device", ib_devices, "--moe-a2a-backend", - "deepep", + "mooncake", "--deepep-mode", "low_latency", "--chunked-prefill-size", @@ -70,6 +70,7 @@ class TestTP(CustomTestCase): self.assertGreater(metrics["accuracy"], 0.60) +@unittest.skip("covered in TestMooncakeWithEPLB") class TestPureDP(TestTP): extra_args = [ "--tp", @@ -90,6 +91,7 @@ class TestHybridDPTP(TestTP): ] +@unittest.skip("covered in TestMooncakeWithEPLB") class TestNoGatherdBuffer(TestTP): extra_args = [ "--tp", @@ -102,6 +104,7 @@ class TestNoGatherdBuffer(TestTP): ] +@unittest.skip("covered in TestMooncakeWithEPLB") class TestTBO(TestTP): extra_args = [ "--tp", @@ -115,7 +118,7 @@ class TestTBO(TestTP): ] -class TestMooncakeWitchEPLB(TestTP): +class TestMooncakeWithEPLB(TestTP): extra_args = [ "--tp", "4", diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index ca0782127..3e1f0237e 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -45,8 +45,7 @@ suites = { ], "per-commit-4-gpu-deepep": [ TestFile("ep/test_deepep_small.py", 531), - # TODO: Add it back after mooncake supports torch 2.9 - # TestFile("ep/test_mooncake_ep_small.py", 450), + TestFile("ep/test_mooncake_ep_small.py", 660), ], "per-commit-8-gpu-h200-deepep": [ TestFile("ep/test_deepep_large.py", 563),