[Bugfix] Fix DeepEP timeout when compiling DeepGeMM in EP+DP+TP (#23185)

Co-authored-by: Byron Hsu <byronhsu@Byrons-MacBook-Pro.local>
Co-authored-by: Cheng Wan <ch-wan@users.noreply.github.com>

python/sglang/compile_deep_gemm.py

@@ -58,17 +58,33 @@ async def warm_up_compile(
     disaggregation_mode: str, tokenizer_manager: TokenizerManager
 ):
     print("\nGenerate warm up request for compiling DeepGEMM...\n")
-    generate_req_input = GenerateReqInput(
-        input_ids=[0, 1, 2, 3],
-        sampling_params={
-            "temperature": 0.0,
-            "max_new_tokens": 8,
-            "ignore_eos": True,
-        },
-    )
+    server_args = tokenizer_manager.server_args
+    dp_size = server_args.dp_size
+    base_ids = [0, 1, 2, 3]
+    sampling_params = {
+        "temperature": 0.0,
+        "max_new_tokens": 8,
+        "ignore_eos": True,
+    }
+
     if disaggregation_mode != "null":
-        generate_req_input.bootstrap_room = 0
-        generate_req_input.bootstrap_host = FAKE_BOOTSTRAP_HOST
+        input_ids = [list(base_ids) for _ in range(dp_size)]
+        generate_req_input = GenerateReqInput(
+            input_ids=input_ids,
+            sampling_params=sampling_params,
+        )
+        generate_req_input.bootstrap_host = [FAKE_BOOTSTRAP_HOST] * dp_size
+        generate_req_input.bootstrap_room = [
+            i * (2**63 // dp_size) + (i % server_args.tp_size) for i in range(dp_size)
+        ]
+    else:
+        input_ids = (
+            base_ids if dp_size == 1 else [list(base_ids) for _ in range(dp_size)]
+        )
+        generate_req_input = GenerateReqInput(
+            input_ids=input_ids,
+            sampling_params=sampling_params,
+        )
 
     await tokenizer_manager.generate_request(generate_req_input, None).__anext__()
 
@@ -104,17 +120,27 @@ def launch_server_process_and_send_one_request(
             if response.status_code == 200:
                 # Rank-0 node send a request to sync with other node and then return.
                 if server_args.node_rank == 0:
+                    dp_size = server_args.dp_size
+                    base_ids = [0, 1, 2, 3]
                     payload = {
-                        "input_ids": [0, 1, 2, 3],
                         "sampling_params": {
                             "max_new_tokens": 8,
                             "temperature": 0,
                         },
                     }
-                    # In PD mode, include fake bootstrap fields so workers don't assert
                     if server_args.disaggregation_mode != "null":
-                        payload["bootstrap_host"] = FAKE_BOOTSTRAP_HOST
-                        payload["bootstrap_room"] = 0
+                        payload["input_ids"] = [list(base_ids) for _ in range(dp_size)]
+                        payload["bootstrap_host"] = [FAKE_BOOTSTRAP_HOST] * dp_size
+                        payload["bootstrap_room"] = [
+                            i * (2**63 // dp_size) + (i % server_args.tp_size)
+                            for i in range(dp_size)
+                        ]
+                    else:
+                        payload["input_ids"] = (
+                            base_ids
+                            if dp_size == 1
+                            else [list(base_ids) for _ in range(dp_size)]
+                        )
 
                     response = requests.post(
                         f"{base_url}/generate",

python/sglang/srt/layers/moe/token_dispatcher/deepep.py

@@ -5,6 +5,7 @@ from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union
 
+from sglang.srt.distributed.parallel_state import get_tp_group
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers import deep_gemm_wrapper
@@ -60,8 +61,16 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
 logger = logging.getLogger(__name__)
 
 
-class DeepEPPDispatchHooks(DispatcherBaseHooks):
+def _deepep_precompile_tp_barrier() -> None:
+    # DeepEP's all-to-all operation has a much shorter timeout compared to torch.distributed,
+    # so if different ranks compile at different speeds, it may quickly trigger a timeout.
+    # To avoid this, we use torch.distributed's barrier during the compile stage.
+    # We apply this barrier only in the compile stage to prevent extra all-reduce overhead at runtime.
+    if envs.SGLANG_IN_DEEPGEMM_PRECOMPILE_STAGE.get():
+        get_tp_group().barrier()
 
+
+class DeepEPPDispatchHooks(DispatcherBaseHooks):
     def __call__(self, dispatcher: BaseDispatcher):
         for hook_fun in self.hook_dict.values():
             hook_fun(dispatcher)
@@ -448,6 +457,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
         # However, doing this would incur an unknown synchronization error, but keeping
         # `handle` as a member variable works.
 
+        _deepep_precompile_tp_barrier()
         (
             recv_x,
             recv_topk_ids,
@@ -508,6 +518,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
 
     def _combine_core(self, x: torch.Tensor, previous_event):
         buffer = self._get_buffer()
+        _deepep_precompile_tp_barrier()
         combined_x, _, event = buffer.combine(
             x,
             self.handle,
@@ -613,6 +624,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
             use_fp8 = True
 
         buffer = self._get_buffer()
+        _deepep_precompile_tp_barrier()
         packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = (
             buffer.low_latency_dispatch(
                 hidden_states,
@@ -695,6 +707,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
             overlap_args_dict = {}
 
         with ctx:
+            _deepep_precompile_tp_barrier()
             combined_hidden_states, event, hook = buffer.low_latency_combine(
                 x=hidden_states,
                 topk_idx=topk_ids,