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

Co-authored-by: Mick <mickjagger19@icloud.com>

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