[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