From 5babb902a9587e1d96729b9b11d9a9541272499b Mon Sep 17 00:00:00 2001
From: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
Date: Tue, 9 Jun 2026 10:31:00 -0700
Subject: [PATCH] [AMD] fix: handle per-frame 4D shift in native scale-shift
 kernel (#27581)

---
 .../jit_kernel/diffusion/triton/scale_shift.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

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,