DP Attention with Auto DeepEP Dispatch (#7222)

python/sglang/srt/two_batch_overlap.py

@@ -13,7 +13,7 @@ from sglang.srt.layers.communicator import (
 )
 from sglang.srt.layers.moe.ep_moe.token_dispatcher import DeepEPDispatcher
 from sglang.srt.layers.quantization import deep_gemm_wrapper
-from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.managers.schedule_batch import ScheduleBatch, global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.operations import execute_operations, execute_overlapped_operations
 from sglang.srt.operations_strategy import OperationsStrategy
@@ -272,7 +272,11 @@ class TboCudaGraphRunnerPlugin:
 
 class TboDPAttentionPreparer:
     def prepare_all_gather(
-        self, local_batch, deepep_mode, enable_deepep_moe, enable_two_batch_overlap
+        self,
+        local_batch: ScheduleBatch,
+        deepep_mode: DeepEPMode,
+        enable_deepep_moe: bool,
+        enable_two_batch_overlap: bool,
     ):
         self.enable_two_batch_overlap = enable_two_batch_overlap
 
@@ -294,7 +298,7 @@ class TboDPAttentionPreparer:
                 extend_lens=local_batch.extend_lens,
                 token_num_per_seq=token_num_per_seq,
             )
-            resolved_deepep_mode = deepep_mode.resolve(local_batch.forward_mode)
+            resolved_deepep_mode = deepep_mode.resolve(local_batch.is_extend_in_batch)
             local_can_run_tbo = (self.local_tbo_split_seq_index is not None) and not (
                 (
                     local_batch.forward_mode.is_extend()