chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Optimization: Reduce the number of D2H operations (#19424)

Browse Source 
Signed-off-by: wangfakang <fakangwang@gmail.com>
This commit is contained in:
sky
2026-03-04 16:32:42 +08:00
committed by GitHub
parent f07d668ba1
commit 17119a697d
2 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/batch_overlap/two_batch_overlap.py
View File

@@ -416,8 +416,10 @@ class TboDPAttentionPreparer:
return local_can_run_tbo, local_forward_mode
def compute_output(self, partial_global_info):
local_can_run_tbo_aggregated = min(partial_global_info[:, 0].tolist())
forward_modes = partial_global_info[:, 1].tolist()
# 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()
global_forward_mode, forward_mode_agree = self._compute_global_forward_mode(
forward_modes

6
python/sglang/srt/managers/scheduler_dp_attn_mixin.py
View File

@@ -94,8 +94,10 @@ class MLPSyncBatchInfo:
tp0_info = global_info_tensor[:, 0, :]
self.tp0_info = tp0_info
self.global_num_tokens = tp0_info[:, 0].tolist()
self.global_num_tokens_for_logprob = tp0_info[:, 1].tolist()
# 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_cuda_graph = bool(tp0_info[:, 2].min().item())
self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
if _ENABLE_METRICS_DP_ATTENTION: