[Diffusion] Opt qwen-image-edit with fuse_residual_layernorm_scale_shift_gate_select01_kernel (#20395)
Co-authored-by: Yihan Chen <yingluosanqian@gmail.com>
This commit is contained in:
Xiaoyu Zhang
@@ -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)
|
||||
@@ -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
|
||||
|
||||
|
||||
219
python/sglang/jit_kernel/tests/test_qwen_image_modulation.py
Normal file
219
python/sglang/jit_kernel/tests/test_qwen_image_modulation.py
Normal file
@@ -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"])
|
||||
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user