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

[diffusion] improve: improve fuse_scale_shift_kernel with non-blocking op (#18710)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Qingfu Wen
2026-02-25 21:04:20 +08:00
committed by GitHub
co-authored by Mick
parent c144e55462
commit 59b9d1e86d
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
+4 -1
View File
@@ -297,7 +297,10 @@ def fuse_scale_shift_kernel(
# If both scalars and both zero, copy fast-path # If both scalars and both zero, copy fast-path
if need_scale_scalar and need_shift_scalar: if need_scale_scalar and need_shift_scalar:
if (scale_blc.abs().max() == 0) and (shift_blc.abs().max() == 0): if not (
scale_blc.any().to("cpu", non_blocking=True)
or shift_blc.any().to("cpu", non_blocking=True)
):
output.copy_(x) output.copy_(x)
return output return output