adapt to pure tp

Author: RunkaiTao

python/sglang/srt/managers/scheduler.py

@@ -1180,7 +1180,7 @@ def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
                 ret = None
 
         # Handle DP attention
-        if self.server_args.enable_dp_attention:
+        if self.server_args.enable_dp_attention or self.dp_size == 1:
             ret, _ = self.prepare_dp_attn_batch(ret)
 
         return ret

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -245,8 +245,7 @@ def __init__(self, model_runner: ModelRunner):
                 )
             else:
                 self.encoder_lens = None
-
-            if self.enable_dp_attention:
+            if self.enable_dp_attention or self.dp_size == 1:
                 self.gathered_buffer = torch.zeros(
                     (
                         self.max_bs * self.dp_size * self.num_tokens_per_bs,
@@ -288,7 +287,7 @@ def model_capture_mode(self):
             self.model_runner.token_to_kv_pool.capture_mode = False
 
     def can_run(self, forward_batch: ForwardBatch):
-        if self.enable_dp_attention:
+        if self.enable_dp_attention or self.dp_size == 1:
             total_global_tokens = sum(forward_batch.global_num_tokens_cpu)
 
             is_bs_supported = forward_batch.can_run_dp_cuda_graph and (
@@ -369,7 +368,7 @@ def capture_one_batch_size(self, bs: int, forward: Callable):
             encoder_lens = None
         mrope_positions = self.mrope_positions[:, :bs]
 
-        if self.enable_dp_attention:
+        if self.enable_dp_attention or self.dp_size == 1:
             self.global_num_tokens_gpu.copy_(
                 torch.tensor(
                     [
@@ -471,7 +470,7 @@ def replay_prepare(self, forward_batch: ForwardBatch):
         raw_num_token = raw_bs * self.num_tokens_per_bs
 
         # Pad
-        if self.enable_dp_attention:
+        if self.enable_dp_attention or self.dp_size == 1:
             index = bisect.bisect_left(
                 self.capture_bs, sum(forward_batch.global_num_tokens_cpu)
             )
@@ -497,7 +496,7 @@ def replay_prepare(self, forward_batch: ForwardBatch):
             self.encoder_lens[:raw_bs].copy_(forward_batch.encoder_lens)
         if forward_batch.mrope_positions is not None:
             self.mrope_positions[:, :raw_bs].copy_(forward_batch.mrope_positions)
-        if self.enable_dp_attention:
+        if self.enable_dp_attention or self.dp_size == 1:
             self.global_num_tokens_gpu.copy_(forward_batch.global_num_tokens_gpu)
 
         if hasattr(forward_batch.spec_info, "hidden_states"):

python/sglang/srt/model_executor/model_runner.py

@@ -280,10 +280,7 @@ def model_specific_adjustment(self):
 
         if server_args.enable_deepep_moe:
             logger.info("DeepEP is turned on.")
-            assert (
-                server_args.enable_dp_attention == True
-            ), "Currently DeepEP is bind to Attention DP. Set '--enable-dp-attention --enable-deepep-moe'"
-
+            
     def init_torch_distributed(self):
         logger.info("Init torch distributed begin.")