From fe4bc8ebd5210817bcd1dcd80789e1158bbd5739 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Sat, 28 Feb 2026 01:52:57 +0800 Subject: [PATCH] [diffusion] fix: MulAdd 4D path (shift indexing) (#18673) Signed-off-by: Xiaodong Ye --- .../sglang/jit_kernel/diffusion/triton/scale_shift.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index 0f1dc2048..b3a27b64c 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -40,14 +40,15 @@ def _fused_scale_shift_4d_kernel( norm_ptrs = normalized_ptr + row_base + col_offsets out_ptrs = output_ptr + row_base + col_offsets - # Pointers for scale and shift for 4D + # Pointers for scale (per-frame) and shift (per-token) b_idx = pid_row // seq_len t_idx = pid_row % seq_len frame_idx_in_batch = t_idx // frame_seqlen scale_row_idx = b_idx * num_frames + frame_idx_in_batch scale_ptrs = scale_ptr + scale_row_idx * inner_dim + col_offsets - shift_ptrs = shift_ptr + scale_row_idx * inner_dim + col_offsets + # shift is per-token [B*L, C], indexed by pid_row directly + shift_ptrs = shift_ptr + pid_row * inner_dim + col_offsets normalized = tl.load(norm_ptrs, mask=mask, other=0.0) scale = tl.load(scale_ptrs, mask=mask, other=0.0) @@ -241,9 +242,10 @@ def fuse_scale_shift_kernel( ), "seq_len must be divisible by num_frames for 4D scale/shift" frame_seqlen = L // num_frames - # Compact [B, F, C] without the singleton dim into [B*F, C] + # Compact scale [B, F, 1, C] -> [B*F, C] (per-frame) scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous() - shift_reshaped = shift.squeeze(2).reshape(-1, C).contiguous() + # 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,