diff --git a/python/sglang/srt/state_capturer/routed_experts.py b/python/sglang/srt/state_capturer/routed_experts.py index fb9a56067..eab516cb6 100644 --- a/python/sglang/srt/state_capturer/routed_experts.py +++ b/python/sglang/srt/state_capturer/routed_experts.py @@ -6,10 +6,13 @@ import torch from sglang.srt.configs.model_config import ModelConfig from sglang.srt.layers.dp_attention import ( + attn_tp_all_gather_into_tensor, get_attention_dp_rank, + get_attention_tp_size, get_dp_local_info, is_dp_attention_enabled, ) +from sglang.srt.layers.moe import get_moe_a2a_backend from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import get_global_server_args from sglang.srt.state_capturer.base import BaseTopkCapturer @@ -75,13 +78,40 @@ class RoutedExpertsCapturer(BaseTopkCapturer): device_topk_size=topk_size + num_fused_shared_experts, ) + # DeepEP a2a path: each attn-TP rank only sees its scattered slice of + # topk_ids. All-gather across attn-TP at capture time so device_cache + # holds the full batch and the existing _get_local_slice / D2H sync + # paths work unchanged. Pre-allocate the gather target. + if get_moe_a2a_backend().is_deepep(): + attn_tp_size = get_attention_tp_size() if is_dp_attention_enabled() else 1 + self.gather_buffer = torch.empty( + ( + self.device_cache.buffer.shape[0] * attn_tp_size, + self.device_cache.buffer.shape[2], + ), + dtype=torch.int32, + device=device, + ) + + def capture(self, layer_id: int, topk_indices: torch.Tensor): + if get_moe_a2a_backend().is_deepep(): + local_topk = topk_indices + topk_indices = self.gather_buffer[ + : local_topk.size(0) * get_attention_tp_size() + ] + attn_tp_all_gather_into_tensor(topk_indices, local_topk) + super().capture(layer_id, topk_indices) + def _get_local_slice( self, forward_batch: ForwardBatch, can_run_graph: bool, cuda_graph_batch: Optional[int], ) -> torch.Tensor: - if is_dp_attention_enabled(): + # Under DeepEP, capture() already attn_tp_all_gathered into the head of + # the per-rank buffer, so the local DP rank's data lives at [0:N_local] + # rather than at the global [start_pos:end_pos] offset. + if is_dp_attention_enabled() and not get_moe_a2a_backend().is_deepep(): local_start_pos, local_num_tokens = get_dp_local_info(forward_batch) if can_run_graph: local_start_pos = get_attention_dp_rank() * cuda_graph_batch diff --git a/test/registered/rl/test_return_routed_experts.py b/test/registered/rl/test_return_routed_experts.py index a7be49e61..8caa66795 100644 --- a/test/registered/rl/test_return_routed_experts.py +++ b/test/registered/rl/test_return_routed_experts.py @@ -13,21 +13,20 @@ from sglang.srt.state_capturer.routed_experts import ( extract_routed_experts_from_meta_info, ) from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci +from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( - DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, CustomTestCase, popen_launch_server, ) -register_cuda_ci(est_time=185, suite="stage-b-test-2-gpu-large") -register_amd_ci( - est_time=200, - suite="stage-b-test-2-gpu-large-amd", - disabled="TP=2 DP=2 routed expert mismatch >15% on AMD; needs TP/DP tuning + concurrency reduction", -) +register_cuda_ci(est_time=400, suite="stage-c-test-4-gpu-h100") + +# FP8 variant of Qwen3-30B-A3B: required because DeepEP normal/LL fast paths in +# ep_moe/layer.py only run for {Fp8Config (via deep_gemm), W4AFp8Config, aiter, +# NPU, modelopt_fp4+cutedsl}. Bf16 hits an `assert False, "deprecated"` today. +MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8" SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered" SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json" @@ -35,34 +34,43 @@ logger = logging.getLogger(__name__) class TestReturnRoutedExperts(CustomTestCase): - # modified from test_hicache.py + """End-to-end check that --enable-return-routed-experts stays correct + under DeepEP a2a + attn_tp_size > 1, across overlap/cuda-graph/radix + optimisations. + + Both servers run ``--tp 4 --dp 2 --enable-dp-attention --moe-a2a-backend + deepep`` so attn_tp_size=2 and the all-gather hot path in + RoutedExpertsCapturer.capture is hit on every step. Baseline disables + overlap/cuda-graph/radix to give a deterministic ground truth; reference + leaves them on. If the gather were skipping a rank or racing against the + forward stream, the captured topk_ids would diverge between the two. + """ + @classmethod def setUpClass(cls): - - cls.baseline_args = [ + common = [ "--enable-return-routed-experts", "--enable-deterministic-inference", + "--tp", + 4, + "--dp", + 2, + "--enable-dp-attention", + "--moe-a2a-backend", + "deepep", + # Force normal-mode dispatch: deepep auto routes decode through + # low_latency mode whose buffer (num_max_dispatch_tokens_per_rank) + # is undersized for cuda graph capture at default --cuda-graph-max-bs. + "--deepep-mode", + "normal", + ] + cls.baseline_args = common + [ "--disable-overlap-schedule", "--disable-cuda-graph", "--disable-radix-cache", - "--tp", - 2, - "--dp", - 2, - "--enable-dp-attention", ] - cls.reference_args = [ - "--enable-return-routed-experts", - "--enable-deterministic-inference", - "--tp", - 2, - "--dp", - 2, - "--enable-dp-attention", - ] - cls.sampling_args = { - "temperature": 0, - } + cls.reference_args = common + cls.sampling_args = {"temperature": 0} # prepare ShareGPT dataset dataset_path = download_and_cache_hf_file(SHAREGPT_REPO_ID, SHAREGPT_FILENAME) with open(dataset_path) as f: @@ -147,7 +155,7 @@ class TestReturnRoutedExperts(CustomTestCase): other_args, ): process = popen_launch_server( - DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST, + MODEL_PATH, DEFAULT_URL_FOR_TEST, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, other_args=other_args,