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

[AMD] fix: handle per-frame 4D shift in native scale-shift kernel (#27581)

Browse Source
This commit is contained in:
Michael
2026-06-09 10:31:00 -07:00
committed by GitHub
parent aa18a68ac5
commit 5babb902a9
1 changed files with 16 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
+16 -2
View File
@@ -363,8 +363,22 @@ def fuse_scale_shift_kernel(
# Compact scale [B, F, 1, C] -> [B*F, C] (per-frame) # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame)
scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous() scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous()
# shift is per-token [B, L, C] -> [B*L, C] if shift.dim() == 4 and current_platform.is_hip():
shift_reshaped = shift.reshape(rows, C).contiguous() # 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]( _fused_scale_shift_4d_kernel[grid](
output_2d, output_2d,