[perf] Collapse the DP attention scheduler sync to a single D2H copy (#34338)

Co-authored-by: Alex Nails <alex.nails@radixark.ai>

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -447,11 +447,9 @@ class TboDPAttentionPreparer:
 
         return local_can_run_tbo, local_forward_mode
 
-    def compute_output(self, partial_global_info):
-        # Perform only one Device-to-Host (D2H) memory copy
-        cpu_data = partial_global_info[:, :2].cpu()
-        local_can_run_tbo_aggregated = min(cpu_data[:, 0].tolist())
-        forward_modes = cpu_data[:, 1].tolist()
+    def compute_output(self, partial_global_info_cpu):
+        local_can_run_tbo_aggregated = min(partial_global_info_cpu[:, 0].tolist())
+        forward_modes = partial_global_info_cpu[:, 1].tolist()
 
         global_forward_mode, forward_mode_agree = self._compute_global_forward_mode(
             forward_modes

python/sglang/srt/managers/scheduler_components/dp_attn.py

@@ -91,7 +91,7 @@ class MLPSyncBatchInfo:
     local_forward_mode: int
 
     # some gathered elements
-    tp0_info: torch.Tensor = None
+    tp0_info_cpu: torch.Tensor = None
     global_num_tokens: list[int] = None
     global_num_tokens_for_logprob: list[int] = None
     tbo_split_seq_index: torch.Tensor = None
@@ -179,17 +179,21 @@ class MLPSyncBatchInfo:
             )
         tp_info[tp_active_ranks[:num_ranks_in_tp_info] == 0] = fallback_tensor
 
-        tp0_info = global_info_tensor[:, 0, :]
-        self.tp0_info = tp0_info
-        # Perform only one Device-to-Host (D2H) memory copy
-        cpu_data = tp0_info[:, :2].cpu()
-        self.global_num_tokens = cpu_data[:, 0].tolist()
-        self.global_num_tokens_for_logprob = cpu_data[:, 1].tolist()
-        self.can_run_decode_cuda_graph = bool(tp0_info[:, 2].min().item())
-        self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
-        self.can_run_prefill_cuda_graph = bool(tp0_info[:, 6].min().item())
+        # One D2H for every field: each `.item()` / `.tolist()` on a device
+        # tensor is its own stream sync. Copy the whole tensor, not the
+        # `[:, 0, :]` slice -- that slice is non-contiguous once
+        # attn_tp * attn_cp > 1, adding a gather kernel inside the wait.
+        tp0_info_cpu = global_info_tensor.cpu()[:, 0, :]
+        self.tp0_info_cpu = tp0_info_cpu
+        self.global_num_tokens = tp0_info_cpu[:, 0].tolist()
+        self.global_num_tokens_for_logprob = tp0_info_cpu[:, 1].tolist()
+        self.can_run_decode_cuda_graph = bool(tp0_info_cpu[:, 2].min())
+        self.is_extend_in_batch = bool(tp0_info_cpu[:, 3].max())
+        self.can_run_prefill_cuda_graph = bool(tp0_info_cpu[:, 6].min())
         if _ENABLE_METRICS_DP_ATTENTION:
-            self.dp_cooperation_info = DPCooperationInfo.create(tp0_info[:, 5].tolist())
+            self.dp_cooperation_info = DPCooperationInfo.create(
+                tp0_info_cpu[:, 5].tolist()
+            )
 
 
 def _update_gather_batch(
@@ -344,7 +348,7 @@ def prepare_mlp_sync_batch_raw(
 
         mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = (
             tbo_preparer.compute_output(
-                mlp_sync_info.tp0_info[:, 4:6],
+                mlp_sync_info.tp0_info_cpu[:, 4:6],
             )
         )
 
@@ -373,7 +377,7 @@ def prepare_mlp_sync_batch_raw(
     if local_batch is not None and not skip_all_gather:
         local_batch.recv_skipper_forward_mode = (
             SchedulerRecvSkipper.derive_forward_mode(
-                mlp_sync_info.tp0_info[:, 5].tolist()
+                mlp_sync_info.tp0_info_cpu[:, 5].tolist()
             )
         )