diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index cd593b340..00b5e77fa 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -363,8 +363,22 @@ def fuse_scale_shift_kernel( # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame) scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous() - # shift is per-token [B, L, C] -> [B*L, C] - shift_reshaped = shift.reshape(rows, C).contiguous() + if shift.dim() == 4 and current_platform.is_hip(): + # ROCm has no fused CUTLASS scale-shift kernel, so this native path + # handles the causal Wan / LingBot output AdaLN, which passes a + # per-frame shift [B, F, 1, C]. Broadcast it across each frame's + # tokens to per-token [B, L, C] before flattening to [B*L, C], + # matching the per-token indexing in _fused_scale_shift_4d_kernel + # (the CUDA fused path accepts [B, F, 1, C] shift and broadcasts it + # per-frame). + shift_reshaped = ( + shift.expand(B, num_frames, frame_seqlen, C) + .reshape(rows, C) + .contiguous() + ) + else: + # shift is per-token [B, L, C] -> [B*L, C] + shift_reshaped = shift.reshape(rows, C).contiguous() _fused_scale_shift_4d_kernel[grid]( output_2d,