minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Diffusion] Avoid GPU syncs in UniPC scheduler (#27440)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-06-06 22:01:41 +08:00
committed by GitHub
parent 4e14b50c48
commit 7c6f9542c7
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py
+6 -6
View File
@@ -853,8 +853,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
rks.append(rk) rks.append(rk)
D1s.append((mi - m0) / rk) D1s.append((mi - m0) / rk)
rks.append(1.0) rks.append(torch.ones((), dtype=h.dtype, device=h.device))
rks = torch.tensor(rks, device=device) rks = torch.stack(rks).to(device=device)
R = [] R = []
b = [] b = []
@@ -879,7 +879,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
h_phi_k = h_phi_k / hh - 1 / factorial_i h_phi_k = h_phi_k / hh - 1 / factorial_i
R = torch.stack(R) R = torch.stack(R)
b = torch.tensor(b, device=device) b = torch.stack(b).to(device=device)
if len(D1s) > 0: if len(D1s) > 0:
D1s = torch.stack(D1s, dim=1) # (B, K) D1s = torch.stack(D1s, dim=1) # (B, K)
@@ -991,8 +991,8 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
rks.append(rk) rks.append(rk)
D1s.append((mi - m0) / rk) D1s.append((mi - m0) / rk)
rks.append(1.0) rks.append(torch.ones((), dtype=h.dtype, device=h.device))
rks = torch.tensor(rks, device=device) rks = torch.stack(rks).to(device=device)
R = [] R = []
b = [] b = []
@@ -1017,7 +1017,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
h_phi_k = h_phi_k / hh - 1 / factorial_i h_phi_k = h_phi_k / hh - 1 / factorial_i
R = torch.stack(R) R = torch.stack(R)
b = torch.tensor(b, device=device) b = torch.stack(b).to(device=device)
if len(D1s) > 0: if len(D1s) > 0:
D1s = torch.stack(D1s, dim=1) D1s = torch.stack(D1s, dim=1)