[Diffusion] Opt qwen-image-edit with fuse_residual_layernorm_scale_shift_gate_select01_kernel (#20395)

Co-authored-by: Yihan Chen <yingluosanqian@gmail.com>
2026-03-13 13:15:22 +08:00
parent 197f807134
commit e00328d1e5
5 changed files with 858 additions and 55 deletions
--- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py
+++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py
@@ -5,6 +5,234 @@ import triton.language as tl  # type: ignore
 from sglang.multimodal_gen.runtime.platforms import current_platform


+@triton.jit
+def _fused_layernorm_scale_shift_gate_select01_kernel(
+    output_ptr,
+    gate_out_ptr,
+    x_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale0_ptr,
+    shift0_ptr,
+    gate0_ptr,
+    scale1_ptr,
+    shift1_ptr,
+    gate1_ptr,
+    index_ptr,
+    inner_dim,
+    seq_len,
+    stride_x_row,
+    stride_out_row,
+    stride_go_row,
+    stride_w,
+    stride_b,
+    stride_s0_b,
+    stride_s0_c,
+    stride_sh0_b,
+    stride_sh0_c,
+    stride_g0_b,
+    stride_g0_c,
+    stride_s1_b,
+    stride_s1_c,
+    stride_sh1_b,
+    stride_sh1_c,
+    stride_g1_b,
+    stride_g1_c,
+    stride_i_b,
+    stride_i_l,
+    eps,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < inner_dim
+
+    x_row_ptr = x_ptr + row * stride_x_row
+    out_row_ptr = output_ptr + row * stride_out_row
+    gate_row_ptr = gate_out_ptr + row * stride_go_row
+
+    x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    mean = tl.sum(x, axis=0) / inner_dim
+    xbar = tl.where(mask, x - mean, 0.0)
+    var = tl.sum(xbar * xbar, axis=0) / inner_dim
+    rstd = tl.rsqrt(var + eps)
+    x_hat = (x - mean) * rstd
+
+    if HAS_WEIGHT:
+        w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+        x_hat = x_hat * w
+    if HAS_BIAS:
+        b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+        x_hat = x_hat + b
+
+    batch_idx = row // seq_len
+    seq_idx = row % seq_len
+    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+    scale0 = tl.load(
+        scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    shift0 = tl.load(
+        shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    gate0 = tl.load(
+        gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c,
+        mask=mask,
+        other=0.0,
+    )
+
+    scale1 = tl.load(
+        scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    shift1 = tl.load(
+        shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    gate1 = tl.load(
+        gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c,
+        mask=mask,
+        other=0.0,
+    )
+
+    scale = tl.where(idx, scale1, scale0)
+    shift = tl.where(idx, shift1, shift0)
+    gate = tl.where(idx, gate1, gate0)
+    y = x_hat * (1.0 + scale) + shift
+
+    tl.store(out_row_ptr + cols, y, mask=mask)
+    tl.store(gate_row_ptr + cols, gate, mask=mask)
+
+
+@triton.jit
+def _fused_residual_layernorm_scale_shift_gate_select01_kernel(
+    output_ptr,
+    residual_out_ptr,
+    gate_out_ptr,
+    x_ptr,
+    residual_ptr,
+    residual_gate_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale0_ptr,
+    shift0_ptr,
+    gate0_ptr,
+    scale1_ptr,
+    shift1_ptr,
+    gate1_ptr,
+    index_ptr,
+    inner_dim,
+    seq_len,
+    stride_x_row,
+    stride_res_row,
+    stride_rg_row,
+    stride_out_row,
+    stride_res_out_row,
+    stride_go_row,
+    stride_w,
+    stride_b,
+    stride_s0_b,
+    stride_s0_c,
+    stride_sh0_b,
+    stride_sh0_c,
+    stride_g0_b,
+    stride_g0_c,
+    stride_s1_b,
+    stride_s1_c,
+    stride_sh1_b,
+    stride_sh1_c,
+    stride_g1_b,
+    stride_g1_c,
+    stride_i_b,
+    stride_i_l,
+    eps,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < inner_dim
+
+    x_row_ptr = x_ptr + row * stride_x_row
+    res_row_ptr = residual_ptr + row * stride_res_row
+    rg_row_ptr = residual_gate_ptr + row * stride_rg_row
+    out_row_ptr = output_ptr + row * stride_out_row
+    res_out_row_ptr = residual_out_ptr + row * stride_res_out_row
+    gate_row_ptr = gate_out_ptr + row * stride_go_row
+
+    x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual = tl.load(res_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual_gate = tl.load(rg_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual_out = residual + residual_gate * x
+    tl.store(res_out_row_ptr + cols, residual_out, mask=mask)
+
+    mean = tl.sum(residual_out, axis=0) / inner_dim
+    xbar = tl.where(mask, residual_out - mean, 0.0)
+    var = tl.sum(xbar * xbar, axis=0) / inner_dim
+    rstd = tl.rsqrt(var + eps)
+    x_hat = (residual_out - mean) * rstd
+
+    if HAS_WEIGHT:
+        w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+        x_hat = x_hat * w
+    if HAS_BIAS:
+        b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+        x_hat = x_hat + b
+
+    batch_idx = row // seq_len
+    seq_idx = row % seq_len
+    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+    scale0 = tl.load(
+        scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    shift0 = tl.load(
+        shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    gate0 = tl.load(
+        gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c,
+        mask=mask,
+        other=0.0,
+    )
+
+    scale1 = tl.load(
+        scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    shift1 = tl.load(
+        shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c,
+        mask=mask,
+        other=0.0,
+    ).to(tl.float32)
+    gate1 = tl.load(
+        gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c,
+        mask=mask,
+        other=0.0,
+    )
+
+    scale = tl.where(idx, scale1, scale0)
+    shift = tl.where(idx, shift1, shift0)
+    gate = tl.where(idx, gate1, gate0)
+    y = x_hat * (1.0 + scale) + shift
+
+    tl.store(out_row_ptr + cols, y, mask=mask)
+    tl.store(gate_row_ptr + cols, gate, mask=mask)
+
+
@triton.autotune(
    configs=[
        triton.Config({"BLOCK_N": 64}, num_warps=2),
@@ -407,6 +635,202 @@ def fuse_scale_shift_gate_select01_kernel(
    return output, gate_out


+def fuse_layernorm_scale_shift_gate_select01_kernel(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    scale0: torch.Tensor,
+    shift0: torch.Tensor,
+    gate0: torch.Tensor,
+    scale1: torch.Tensor,
+    shift1: torch.Tensor,
+    gate1: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+):
+    assert x.is_cuda
+    assert x.is_contiguous()
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    gate_out = torch.empty_like(x)
+
+    if (
+        scale0.dim() != 2
+        or shift0.dim() != 2
+        or gate0.dim() != 2
+        or scale1.dim() != 2
+        or shift1.dim() != 2
+        or gate1.dim() != 2
+    ):
+        raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
+    if index.dim() != 2:
+        raise ValueError("index must be 2D [B, L]")
+    if weight is not None and (weight.dim() != 1 or weight.shape[0] != C):
+        raise ValueError("weight must be 1D [C]")
+    if bias is not None and (bias.dim() != 1 or bias.shape[0] != C):
+        raise ValueError("bias must be 1D [C]")
+
+    x_2d = x.view(B * L, C)
+    output_2d = output.view(B * L, C)
+    gate_out_2d = gate_out.view(B * L, C)
+    weight = weight.contiguous() if weight is not None else x_2d
+    bias = bias.contiguous() if bias is not None else x_2d
+
+    MAX_FUSED_SIZE = 65536 // x_2d.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C))
+    if C > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+
+    grid = (B * L,)
+    _fused_layernorm_scale_shift_gate_select01_kernel[grid](
+        output_2d,
+        gate_out_2d,
+        x_2d,
+        weight,
+        bias,
+        scale0.contiguous(),
+        shift0.contiguous(),
+        gate0.contiguous(),
+        scale1.contiguous(),
+        shift1.contiguous(),
+        gate1.contiguous(),
+        index.contiguous(),
+        C,
+        L,
+        x_2d.stride(0),
+        output_2d.stride(0),
+        gate_out_2d.stride(0),
+        weight.stride(0) if weight.dim() == 1 else 0,
+        bias.stride(0) if bias.dim() == 1 else 0,
+        scale0.stride(0),
+        scale0.stride(1),
+        shift0.stride(0),
+        shift0.stride(1),
+        gate0.stride(0),
+        gate0.stride(1),
+        scale1.stride(0),
+        scale1.stride(1),
+        shift1.stride(0),
+        shift1.stride(1),
+        gate1.stride(0),
+        gate1.stride(1),
+        index.stride(0),
+        index.stride(1),
+        eps,
+        HAS_WEIGHT=weight is not x_2d,
+        HAS_BIAS=bias is not x_2d,
+        BLOCK_N=BLOCK_N,
+    )
+    return output, gate_out
+
+
+def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    residual_gate: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    scale0: torch.Tensor,
+    shift0: torch.Tensor,
+    gate0: torch.Tensor,
+    scale1: torch.Tensor,
+    shift1: torch.Tensor,
+    gate1: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+):
+    assert x.is_cuda
+    assert x.is_contiguous()
+    assert residual.is_contiguous()
+    assert residual_gate.is_contiguous()
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    residual_out = torch.empty_like(x)
+    gate_out = torch.empty_like(x)
+
+    if residual.shape != x.shape:
+        raise ValueError("residual must have the same shape as x")
+    if residual_gate.shape != x.shape:
+        raise ValueError("residual_gate must have the same shape as x")
+    if (
+        scale0.dim() != 2
+        or shift0.dim() != 2
+        or gate0.dim() != 2
+        or scale1.dim() != 2
+        or shift1.dim() != 2
+        or gate1.dim() != 2
+    ):
+        raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
+    if index.dim() != 2:
+        raise ValueError("index must be 2D [B, L]")
+    if weight is not None and (weight.dim() != 1 or weight.shape[0] != C):
+        raise ValueError("weight must be 1D [C]")
+    if bias is not None and (bias.dim() != 1 or bias.shape[0] != C):
+        raise ValueError("bias must be 1D [C]")
+
+    x_2d = x.view(B * L, C)
+    residual_2d = residual.view(B * L, C)
+    residual_gate_2d = residual_gate.view(B * L, C)
+    output_2d = output.view(B * L, C)
+    residual_out_2d = residual_out.view(B * L, C)
+    gate_out_2d = gate_out.view(B * L, C)
+    weight = weight.contiguous() if weight is not None else x_2d
+    bias = bias.contiguous() if bias is not None else x_2d
+
+    MAX_FUSED_SIZE = 65536 // x_2d.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C))
+    if C > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+
+    grid = (B * L,)
+    _fused_residual_layernorm_scale_shift_gate_select01_kernel[grid](
+        output_2d,
+        residual_out_2d,
+        gate_out_2d,
+        x_2d,
+        residual_2d,
+        residual_gate_2d,
+        weight,
+        bias,
+        scale0.contiguous(),
+        shift0.contiguous(),
+        gate0.contiguous(),
+        scale1.contiguous(),
+        shift1.contiguous(),
+        gate1.contiguous(),
+        index.contiguous(),
+        C,
+        L,
+        x_2d.stride(0),
+        residual_2d.stride(0),
+        residual_gate_2d.stride(0),
+        output_2d.stride(0),
+        residual_out_2d.stride(0),
+        gate_out_2d.stride(0),
+        weight.stride(0) if weight.dim() == 1 else 0,
+        bias.stride(0) if bias.dim() == 1 else 0,
+        scale0.stride(0),
+        scale0.stride(1),
+        shift0.stride(0),
+        shift0.stride(1),
+        gate0.stride(0),
+        gate0.stride(1),
+        scale1.stride(0),
+        scale1.stride(1),
+        shift1.stride(0),
+        shift1.stride(1),
+        gate1.stride(0),
+        gate1.stride(1),
+        index.stride(0),
+        index.stride(1),
+        eps,
+        HAS_WEIGHT=weight is not x_2d,
+        HAS_BIAS=bias is not x_2d,
+        BLOCK_N=BLOCK_N,
+    )
+    return output, residual_out, gate_out
+
+
 if current_platform.is_npu():
    from .npu_fallback import fuse_scale_shift_native