From 59b9d1e86db0a5ad8d73bd77c9050021bbfa7021 Mon Sep 17 00:00:00 2001 From: Qingfu Wen Date: Wed, 25 Feb 2026 21:04:20 +0800 Subject: [PATCH] [diffusion] improve: improve fuse_scale_shift_kernel with non-blocking op (#18710) Co-authored-by: Mick --- python/sglang/jit_kernel/diffusion/triton/scale_shift.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index 16ff4461e..0f1dc2048 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -297,7 +297,10 @@ def fuse_scale_shift_kernel( # If both scalars and both zero, copy fast-path 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) return output