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

Co-authored-by: Yihan Chen <yingluosanqian@gmail.com>

python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py

@@ -0,0 +1,178 @@
+from typing import Tuple
+
+import torch
+import triton.testing
+
+from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph
+from sglang.jit_kernel.diffusion.triton.norm import norm_infer
+from sglang.jit_kernel.diffusion.triton.scale_shift import (
+    fuse_layernorm_scale_shift_gate_select01_kernel,
+    fuse_residual_layernorm_scale_shift_gate_select01_kernel,
+    fuse_scale_shift_gate_select01_kernel,
+)
+
+if is_in_ci():
+    B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
+else:
+    B_RANGE, S_RANGE, D_RANGE = [1, 2], [128, 512, 2048], [1024, 1536, 3072]
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+EPS = 1e-6
+LINE_VALS = ["split", "fused"]
+LINE_NAMES = ["Split Kernels", "Fused Triton"]
+STYLES = [("red", "-"), ("blue", "--")]
+CONFIG = [(b, s, d) for b in B_RANGE for s in S_RANGE for d in D_RANGE]
+
+
+def _make_common_inputs(batch_size: int, seq_len: int, hidden_size: int):
+    x = torch.randn(batch_size, seq_len, hidden_size, dtype=DTYPE, device=DEVICE)
+    weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
+    bias = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
+    index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int32, device=DEVICE)
+    scale0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
+    shift0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
+    gate0 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
+    scale1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
+    shift1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
+    gate1 = torch.randn(batch_size, hidden_size, dtype=DTYPE, device=DEVICE)
+    return x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["B", "S", "D"],
+        x_vals=CONFIG,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=STYLES,
+        ylabel="us",
+        plot_name="qwen_image_layernorm_scale_shift_gate_select01",
+        args={},
+    )
+)
+def bench_layernorm_scale_shift_gate_select01(
+    B: int, S: int, D: int, provider: str
+) -> Tuple[float, float, float]:
+    x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = (
+        _make_common_inputs(B, S, D)
+    )
+
+    if provider == "split":
+
+        def fn():
+            normalized = norm_infer(
+                x.view(-1, x.shape[-1]),
+                weight,
+                bias,
+                eps=EPS,
+                is_rms_norm=False,
+            ).view_as(x)
+            return fuse_scale_shift_gate_select01_kernel(
+                normalized,
+                scale0=scale0,
+                shift0=shift0,
+                gate0=gate0,
+                scale1=scale1,
+                shift1=shift1,
+                gate1=gate1,
+                index=index,
+            )
+
+    else:
+
+        def fn():
+            return fuse_layernorm_scale_shift_gate_select01_kernel(
+                x,
+                weight=weight,
+                bias=bias,
+                scale0=scale0,
+                shift0=shift0,
+                gate0=gate0,
+                scale1=scale1,
+                shift1=shift1,
+                gate1=gate1,
+                index=index,
+                eps=EPS,
+            )
+
+    return run_benchmark_no_cudagraph(fn)
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["B", "S", "D"],
+        x_vals=CONFIG,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=STYLES,
+        ylabel="us",
+        plot_name="qwen_image_residual_layernorm_scale_shift_gate_select01",
+        args={},
+    )
+)
+def bench_residual_layernorm_scale_shift_gate_select01(
+    B: int, S: int, D: int, provider: str
+) -> Tuple[float, float, float]:
+    x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1 = (
+        _make_common_inputs(B, S, D)
+    )
+    residual = torch.randn_like(x)
+    residual_gate = torch.randn_like(x)
+
+    if provider == "split":
+
+        def fn():
+            residual_out = residual + residual_gate * x
+            normalized = norm_infer(
+                residual_out.view(-1, residual_out.shape[-1]),
+                weight,
+                bias,
+                eps=EPS,
+                is_rms_norm=False,
+            ).view_as(residual_out)
+            return fuse_scale_shift_gate_select01_kernel(
+                normalized,
+                scale0=scale0,
+                shift0=shift0,
+                gate0=gate0,
+                scale1=scale1,
+                shift1=shift1,
+                gate1=gate1,
+                index=index,
+            )
+
+    else:
+
+        def fn():
+            return fuse_residual_layernorm_scale_shift_gate_select01_kernel(
+                x,
+                residual=residual,
+                residual_gate=residual_gate,
+                weight=weight,
+                bias=bias,
+                scale0=scale0,
+                shift0=shift0,
+                gate0=gate0,
+                scale1=scale1,
+                shift1=shift1,
+                gate1=gate1,
+                index=index,
+                eps=EPS,
+            )
+
+    return run_benchmark_no_cudagraph(fn)
+
+
+if __name__ == "__main__":
+    print(f"\n{'=' * 80}")
+    print("Benchmark: qwen_image layernorm + scale_shift_gate_select01")
+    print(f"{'=' * 80}\n")
+    bench_layernorm_scale_shift_gate_select01.run(print_data=True)
+
+    print(f"\n{'=' * 80}")
+    print("Benchmark: qwen_image residual + layernorm + scale_shift_gate_select01")
+    print(f"{'=' * 80}\n")
+    bench_residual_layernorm_scale_shift_gate_select01.run(print_data=True)

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
 

python/sglang/jit_kernel/tests/test_qwen_image_modulation.py

@@ -0,0 +1,219 @@
+import pytest
+import torch
+import triton
+
+from sglang.jit_kernel.diffusion.triton.norm import norm_infer
+from sglang.jit_kernel.diffusion.triton.scale_shift import (
+    fuse_layernorm_scale_shift_gate_select01_kernel,
+    fuse_residual_layernorm_scale_shift_gate_select01_kernel,
+    fuse_scale_shift_gate_select01_kernel,
+)
+from sglang.jit_kernel.utils import get_ci_test_range
+
+DEVICE = "cuda"
+DTYPES = get_ci_test_range(
+    [torch.float16, torch.bfloat16, torch.float32], [torch.float16, torch.bfloat16]
+)
+BATCH_SIZES = get_ci_test_range([1, 2, 4], [1, 2])
+SEQ_LENS = get_ci_test_range([6, 33, 128, 257], [6, 128])
+HIDDEN_SIZES = get_ci_test_range([512, 1024, 1536, 3072], [512, 3072])
+EPS = 1e-6
+
+
+def _tol(dtype: torch.dtype) -> tuple[float, float]:
+    if dtype == torch.float32:
+        return 1e-5, 1e-5
+    return 5e-2, 5e-2
+
+
+def _make_modulation_tensors(batch_size: int, hidden_size: int, dtype: torch.dtype):
+    scale0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
+    shift0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
+    gate0 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
+    scale1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
+    shift1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
+    gate1 = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
+    return scale0, shift0, gate0, scale1, shift1, gate1
+
+
+def _baseline_select01_modulation(
+    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,
+):
+    normalized = norm_infer(
+        x.view(-1, x.shape[-1]),
+        weight,
+        bias,
+        eps=eps,
+        is_rms_norm=False,
+    ).view_as(x)
+    output, gate_out = fuse_scale_shift_gate_select01_kernel(
+        normalized,
+        scale0=scale0,
+        shift0=shift0,
+        gate0=gate0,
+        scale1=scale1,
+        shift1=shift1,
+        gate1=gate1,
+        index=index,
+    )
+    return output, gate_out
+
+
+def _baseline_residual_select01_modulation(
+    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,
+):
+    residual_out = residual + residual_gate * x
+    normalized = norm_infer(
+        residual_out.view(-1, residual_out.shape[-1]),
+        weight,
+        bias,
+        eps=eps,
+        is_rms_norm=False,
+    ).view_as(residual_out)
+    output, gate_out = fuse_scale_shift_gate_select01_kernel(
+        normalized,
+        scale0=scale0,
+        shift0=shift0,
+        gate0=gate0,
+        scale1=scale1,
+        shift1=shift1,
+        gate1=gate1,
+        index=index,
+    )
+    return output, residual_out, gate_out
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("batch_size", BATCH_SIZES)
+@pytest.mark.parametrize("seq_len", SEQ_LENS)
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
+def test_fused_layernorm_scale_shift_gate_select01(
+    dtype, batch_size, seq_len, hidden_size
+):
+    x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
+    weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
+    bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
+    index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
+    scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
+        batch_size, hidden_size, dtype
+    )
+
+    out_ref, gate_ref = _baseline_select01_modulation(
+        x,
+        weight,
+        bias,
+        scale0,
+        shift0,
+        gate0,
+        scale1,
+        shift1,
+        gate1,
+        index,
+        EPS,
+    )
+    out_fused, gate_fused = fuse_layernorm_scale_shift_gate_select01_kernel(
+        x.contiguous(),
+        weight=weight,
+        bias=bias,
+        scale0=scale0,
+        shift0=shift0,
+        gate0=gate0,
+        scale1=scale1,
+        shift1=shift1,
+        gate1=gate1,
+        index=index,
+        eps=EPS,
+    )
+
+    atol, rtol = _tol(dtype)
+    triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
+    triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("batch_size", BATCH_SIZES)
+@pytest.mark.parametrize("seq_len", SEQ_LENS)
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
+def test_fused_residual_layernorm_scale_shift_gate_select01(
+    dtype, batch_size, seq_len, hidden_size
+):
+    x = torch.randn(batch_size, seq_len, hidden_size, device=DEVICE, dtype=dtype)
+    residual = torch.randn_like(x)
+    residual_gate = torch.randn_like(x)
+    weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
+    bias = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
+    index = torch.randint(0, 2, (batch_size, seq_len), device=DEVICE, dtype=torch.int32)
+    scale0, shift0, gate0, scale1, shift1, gate1 = _make_modulation_tensors(
+        batch_size, hidden_size, dtype
+    )
+
+    out_ref, residual_ref, gate_ref = _baseline_residual_select01_modulation(
+        x,
+        residual,
+        residual_gate,
+        weight,
+        bias,
+        scale0,
+        shift0,
+        gate0,
+        scale1,
+        shift1,
+        gate1,
+        index,
+        EPS,
+    )
+    out_fused, residual_fused, gate_fused = (
+        fuse_residual_layernorm_scale_shift_gate_select01_kernel(
+            x.contiguous(),
+            residual=residual.contiguous(),
+            residual_gate=residual_gate.contiguous(),
+            weight=weight,
+            bias=bias,
+            scale0=scale0,
+            shift0=shift0,
+            gate0=gate0,
+            scale1=scale1,
+            shift1=shift1,
+            gate1=gate1,
+            index=index,
+            eps=EPS,
+        )
+    )
+
+    atol, rtol = _tol(dtype)
+    triton.testing.assert_close(out_ref, out_fused, atol=atol, rtol=rtol)
+    triton.testing.assert_close(residual_ref, residual_fused, atol=atol, rtol=rtol)
+    triton.testing.assert_close(gate_ref, gate_fused, atol=atol, rtol=rtol)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -580,14 +580,3 @@ def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor:
     )
     output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight
     return output.to(dtype=src_dtype)
-
-
-# TODO: Workaround, fuse norm with new select01 kernel
-def apply_layernorm_only(x: torch.Tensor, layernorm_scale_shift: LayerNormScaleShift):
-    return norm_infer(
-        x.view(-1, x.shape[-1]),
-        layernorm_scale_shift.norm.weight,
-        layernorm_scale_shift.norm.bias,
-        eps=layernorm_scale_shift.eps,
-        is_rms_norm=False,
-    ).view(x.shape)

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -16,7 +16,8 @@ from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 
 from sglang.jit_kernel.diffusion.triton.scale_shift import (
-    fuse_scale_shift_gate_select01_kernel,
+    fuse_layernorm_scale_shift_gate_select01_kernel,
+    fuse_residual_layernorm_scale_shift_gate_select01_kernel,
 )
 from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
@@ -26,7 +27,6 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
     LayerNormScaleShift,
     RMSNorm,
     ScaleResidualLayerNormScaleShift,
-    apply_layernorm_only,
     apply_qk_norm,
 )
 from sglang.multimodal_gen.runtime.layers.linear import (
@@ -44,15 +44,11 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
-from sglang.multimodal_gen.runtime.platforms import (
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
-    AttentionBackendEnum,
-    current_platform,
-)
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
-_is_cuda = current_platform.is_cuda()
 
 try:
     from nunchaku.models.attention import NunchakuFeedForward  # type: ignore[import]
@@ -808,18 +804,22 @@ class QwenImageTransformerBlock(nn.Module):
                 scale[actual_batch : 2 * actual_batch],
             )
             gate0, gate1 = gate[:actual_batch], gate[actual_batch : 2 * actual_batch]
-            if _is_cuda:
-                if is_scale_residual:
-                    x = gate_x * x + residual_x
-                    residual_out = x
             if not x.is_contiguous():
                 x = x.contiguous()
             if not index.is_contiguous():
                 index = index.contiguous()
-                # TODO: fuse norm with above select01 kernel, workaround now
+            if is_scale_residual:
-                x = apply_layernorm_only(x, norm_module)
+                if not residual_x.is_contiguous():
-                x, gate_result = fuse_scale_shift_gate_select01_kernel(
+                    residual_x = residual_x.contiguous()
+                if not gate_x.is_contiguous():
+                    gate_x = gate_x.contiguous()
+                x, residual_out, gate_result = (
+                    fuse_residual_layernorm_scale_shift_gate_select01_kernel(
                         x,
+                        residual=residual_x,
+                        residual_gate=gate_x,
+                        weight=getattr(norm_module.norm, "weight", None),
+                        bias=getattr(norm_module.norm, "bias", None),
                         scale0=scale0.contiguous(),
                         shift0=shift0.contiguous(),
                         gate0=gate0.contiguous(),
@@ -827,32 +827,25 @@ class QwenImageTransformerBlock(nn.Module):
                         shift1=shift1.contiguous(),
                         gate1=gate1.contiguous(),
                         index=index,
+                        eps=norm_module.eps,
+                    )
                 )
-                if is_scale_residual:
                 return x, residual_out, gate_result
             else:
+                x, gate_result = fuse_layernorm_scale_shift_gate_select01_kernel(
+                    x,
+                    weight=getattr(norm_module.norm, "weight", None),
+                    bias=getattr(norm_module.norm, "bias", None),
+                    scale0=scale0.contiguous(),
+                    shift0=shift0.contiguous(),
+                    gate0=gate0.contiguous(),
+                    scale1=scale1.contiguous(),
+                    shift1=shift1.contiguous(),
+                    gate1=gate1.contiguous(),
+                    index=index,
+                    eps=norm_module.eps,
+                )
                 return x, gate_result
-            else:
-                mask = (index == 0).unsqueeze(-1)
-                shift_result = torch.where(
-                    mask, shift0.unsqueeze(1), shift1.unsqueeze(1)
-                )
-                scale_result = torch.where(
-                    mask, scale0.unsqueeze(1), scale1.unsqueeze(1)
-                )
-                gate_result = torch.where(mask, gate0.unsqueeze(1), gate1.unsqueeze(1))
-                if is_scale_residual:
-                    modulated, residual_out = norm_module(
-                        residual=residual_x,
-                        x=x,
-                        gate=gate_x,
-                        shift=shift_result,
-                        scale=scale_result,
-                    )
-                    return modulated, residual_out, gate_result
-                else:
-                    modulated = norm_module(x=x, shift=shift_result, scale=scale_result)
-                    return modulated, gate_result
         else:
             shift_result = shift.unsqueeze(1)
             scale_result = scale.unsqueeze(1)
@@ -879,7 +872,7 @@ class QwenImageTransformerBlock(nn.Module):
         temb_txt_silu: torch.Tensor,
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
-        modulate_index: Optional[List[int]] = None,
+        modulate_index: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         # Get modulation parameters for both streams
         img_mod_params = self.img_mod[1](temb_img_silu)  # [B, 6*dim]