From 1cff871c67ef98a5bcd09bc282f939c79f146dc9 Mon Sep 17 00:00:00 2001 From: Byron Hsu Date: Sun, 19 Apr 2026 17:36:11 -0700 Subject: [PATCH] [Bugfix] Fix DeepEP timeout when compiling DeepGeMM in EP+DP+TP (#23185) Co-authored-by: Byron Hsu Co-authored-by: Cheng Wan --- python/sglang/compile_deep_gemm.py | 54 ++++++++++++++----- .../srt/layers/moe/token_dispatcher/deepep.py | 15 +++++- 2 files changed, 54 insertions(+), 15 deletions(-) diff --git a/python/sglang/compile_deep_gemm.py b/python/sglang/compile_deep_gemm.py index 77ddbadce..7abc6993b 100644 --- a/python/sglang/compile_deep_gemm.py +++ b/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", diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index a4339808b..b3fb47d3f 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/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,