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