diff --git a/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py b/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py new file mode 100644 index 000000000..d16f9eda7 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index 8be4c11bd..a34db4d18 100644 --- 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 diff --git a/python/sglang/jit_kernel/tests/test_qwen_image_modulation.py b/python/sglang/jit_kernel/tests/test_qwen_image_modulation.py new file mode 100644 index 000000000..ae8f7ef40 --- /dev/null +++ b/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"]) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 120730556..b69040672 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index de2aa35ba..7b10905b9 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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 ( - AttentionBackendEnum, - current_platform, -) +from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum 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,38 @@ 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 - x = apply_layernorm_only(x, norm_module) - x, gate_result = fuse_scale_shift_gate_select01_kernel( + if not x.is_contiguous(): + x = x.contiguous() + if not index.is_contiguous(): + index = index.contiguous() + if is_scale_residual: + if not residual_x.is_contiguous(): + 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(), + scale1=scale1.contiguous(), + shift1=shift1.contiguous(), + gate1=gate1.contiguous(), + index=index, + eps=norm_module.eps, + ) + ) + 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(), @@ -827,32 +843,9 @@ 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: - 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 + return x, 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]