[RL] DeepEP support for --enable-return-routed-experts (#16859)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Junrong Lin <33685709+ocss884@users.noreply.github.com>

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

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_cuda_ci(est_time=400, suite="stage-c-test-4-gpu-h100")
-register_amd_ci(
+
-    est_time=200,
+# FP8 variant of Qwen3-30B-A3B: required because DeepEP normal/LL fast paths in
-    suite="stage-b-test-2-gpu-large-amd",
+# ep_moe/layer.py only run for {Fp8Config (via deep_gemm), W4AFp8Config, aiter,
-    disabled="TP=2 DP=2 routed expert mismatch >15% on AMD; needs TP/DP tuning + concurrency reduction",
+# 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):
-
+        common = [
-        cls.baseline_args = [
             "--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 = [
+        cls.reference_args = common
-            "--enable-return-routed-experts",
+        cls.sampling_args = {"temperature": 0}
-            "--enable-deterministic-inference",
-            "--tp",
-            2,
-            "--dp",
-            2,
-            "--enable-dp-attention",
-        ]
-        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,