diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 2b52ab8ca..1e89fba14 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index af3182611..c0b40307f 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/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() ) )