chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Diffusion] Clean up diffusion Triton kernels and modernize custom op registration (#21122)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
2026-03-22 22:38:57 +08:00
committed by GitHub
parent 2406ddfdb8
commit c1fe5de69c
9 changed files with 257 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py
View File

@@ -8,7 +8,6 @@ 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.utils import is_in_ci
@@ -21,7 +20,7 @@ DTYPE = torch.bfloat16
DEVICE = "cuda"
EPS = 1e-6
LINE_VALS = ["split", "fused"]
LINE_NAMES = ["Split Kernels", "Fused Triton"]
LINE_NAMES = ["Triton Norm + Torch Select", "Fused Triton"]
STYLES = [("red", "-"), ("blue", "--")]
CONFIG = [(b, s, d) for b in B_RANGE for s in S_RANGE for d in D_RANGE]
@@ -40,6 +39,23 @@ def _make_common_inputs(batch_size: int, seq_len: int, hidden_size: int):
return x, weight, bias, index, scale0, shift0, gate0, scale1, shift1, gate1
def _apply_select01_modulation(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
):
idx = index.bool().unsqueeze(-1)
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
return x * (1 + scale) + shift, gate
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["B", "S", "D"],
@@ -70,15 +86,8 @@ def bench_layernorm_scale_shift_gate_select01(
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,
return _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
else:
@@ -134,15 +143,8 @@ def bench_residual_layernorm_scale_shift_gate_select01(
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,
return _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
else:

21
python/sglang/jit_kernel/diffusion/triton/mps_fallback.py
View File

@@ -94,27 +94,6 @@ def fuse_scale_shift_kernel_native(
return x * (scale_constant + scale) + shift
def fuse_scale_shift_gate_select01_kernel_native(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
block_l: int = 128,
block_c: int = 128,
):
"""Native fallback for fuse_scale_shift_gate_select01_kernel."""
idx = index.unsqueeze(-1).bool()
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
y = x * (1 + scale) + shift
return y, gate
def apply_rotary_embedding_native(
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
) -> torch.Tensor:

407
python/sglang/jit_kernel/diffusion/triton/norm.py
View File

@@ -1,11 +1,12 @@
from typing import Optional
from typing import Optional, Tuple
import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from torch import Tensor
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.srt.utils.custom_op import register_custom_op
# RMSNorm-fp32
@@ -19,7 +20,6 @@ def maybe_contiguous(x):
def triton_autotune_configs():
# Return configs with a valid warp count for the current device
configs = []
# Maximum threads per block is architecture-dependent in theory, but in reality all are 1024
max_threads_per_block = 1024
# Default to warp size 32 if not defined by device
@@ -189,8 +189,8 @@ def _layer_norm_fwd_1pass_kernel(
def _layer_norm_fwd(
x: Tensor,
weight: Tensor,
bias: Tensor,
weight: Optional[Tensor],
bias: Optional[Tensor],
eps: float,
residual: Optional[Tensor] = None,
x1: Optional[Tensor] = None,
@@ -206,9 +206,7 @@ def _layer_norm_fwd(
out: Optional[Tensor] = None,
residual_out: Optional[Tensor] = None,
) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor):
# Need to wrap to handle the case where residual_out is a alias of x, which makes torch.library
# and torch.compile unhappy. Also allocate memory for out and residual_out if they are None
# so that _layer_norm_fwd_impl doesn't have to return them.
# Allocate aliases upfront so the custom op only mutates explicit outputs.
if out is None:
out = torch.empty_like(x, dtype=x.dtype if out_dtype is None else out_dtype)
if residual is not None:
@@ -248,25 +246,40 @@ def _layer_norm_fwd(
return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1
# [2025-04-28] torch.library.triton_op ignores the schema argument, but here we need the schema
# since we're returning a tuple of tensors
def _layer_norm_fwd_impl(
@register_custom_op(
op_name="diffusion_layer_norm_fwd_impl_cuda",
mutates_args=[
"out",
"y1",
"mean",
"rstd",
"residual_out",
"dropout_mask",
"dropout_mask1",
],
)
def _layer_norm_fwd_impl_cuda(
x: Tensor,
weight: Optional[Tensor],
bias: Tensor,
bias: Optional[Tensor],
eps: float,
out: Tensor,
y1: Optional[Tensor],
mean: Optional[Tensor],
rstd: Tensor,
residual: Optional[Tensor] = None,
x1: Optional[Tensor] = None,
weight1: Optional[Tensor] = None,
bias1: Optional[Tensor] = None,
dropout_p: float = 0.0,
residual_out: Optional[Tensor] = None,
rowscale: Optional[Tensor] = None,
seeds: Optional[Tensor] = None,
dropout_mask: Optional[Tensor] = None,
dropout_mask1: Optional[Tensor] = None,
dropout_p: float = 0.0,
zero_centered_weight: bool = False,
is_rms_norm: bool = False,
return_dropout_mask: bool = False,
residual_out: Optional[Tensor] = None,
) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor):
) -> None:
M, N = x.shape
assert x.stride(-1) == 1
if residual is not None:
@@ -296,38 +309,25 @@ def _layer_norm_fwd_impl(
if residual_out is not None:
assert residual_out.shape == x.shape
assert residual_out.stride(-1) == 1
if weight1 is not None:
y1 = torch.empty_like(out)
if y1 is not None:
assert y1.shape == x.shape
assert y1.stride(-1) == 1
else:
y1 = None
mean = (
torch.empty((M,), dtype=torch.float32, device=x.device)
if not is_rms_norm
else None
)
rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
if dropout_p > 0.0:
seeds = torch.randint(
2**32, (M if x1 is None else 2 * M,), device=x.device, dtype=torch.int64
)
else:
seeds = None
if return_dropout_mask and dropout_p > 0.0:
dropout_mask = torch.empty(M, N, device=x.device, dtype=torch.bool)
if x1 is not None:
dropout_mask1 = torch.empty(M, N, device=x.device, dtype=torch.bool)
else:
dropout_mask1 = None
else:
dropout_mask, dropout_mask1 = None, None
if mean is not None:
assert mean.shape == (M,)
assert rstd.shape == (M,)
if seeds is not None:
assert seeds.shape == (M if x1 is None else 2 * M,)
if dropout_mask is not None:
assert dropout_mask.shape == (M, N)
if dropout_mask1 is not None:
assert dropout_mask1.shape == (M, N)
# Less than 64KB per feature: enqueue fused kernel
MAX_FUSED_SIZE = 65536 // x.element_size()
BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
if N > BLOCK_N:
raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
with torch.get_device_module().device(x.device.index):
torch.library.wrap_triton(_layer_norm_fwd_1pass_kernel)[(M,)](
with torch.get_device_module().device(x.device):
_layer_norm_fwd_1pass_kernel[(M,)](
x,
out,
weight if weight is not None else x, # unused when HAS_WEIGHT == False
@@ -369,91 +369,83 @@ def _layer_norm_fwd_impl(
HAS_W1=weight1 is not None,
HAS_B1=bias1 is not None,
)
return y1, mean, rstd, seeds, dropout_mask, dropout_mask1
return None
class LayerNormFn:
@staticmethod
def forward(
def _layer_norm_fwd_impl(
x: Tensor,
weight: Optional[Tensor],
bias: Optional[Tensor],
eps: float,
out: Tensor,
residual: Optional[Tensor] = None,
x1: Optional[Tensor] = None,
weight1: Optional[Tensor] = None,
bias1: Optional[Tensor] = None,
dropout_p: float = 0.0,
rowscale: Optional[Tensor] = None,
zero_centered_weight: bool = False,
is_rms_norm: bool = False,
return_dropout_mask: bool = False,
residual_out: Optional[Tensor] = None,
) -> Tuple[
Optional[Tensor],
Optional[Tensor],
Tensor,
Optional[Tensor],
Optional[Tensor],
Optional[Tensor],
]:
M, N = x.shape
y1 = torch.empty_like(out) if weight1 is not None else None
mean = (
torch.empty((M,), dtype=torch.float32, device=x.device)
if not is_rms_norm
else None
)
rstd = torch.empty((M,), dtype=torch.float32, device=x.device)
seeds = (
torch.randint(
2**32, (M if x1 is None else 2 * M), device=x.device, dtype=torch.int64
)
if dropout_p > 0.0
else None
)
if return_dropout_mask and dropout_p > 0.0:
dropout_mask = torch.empty((M, N), dtype=torch.bool, device=x.device)
dropout_mask1 = (
torch.empty((M, N), dtype=torch.bool, device=x.device)
if x1 is not None
else None
)
else:
dropout_mask = dropout_mask1 = None
_layer_norm_fwd_impl_cuda(
x,
weight,
bias,
residual=None,
x1=None,
weight1=None,
bias1=None,
eps=1e-6,
dropout_p=0.0,
rowscale=None,
prenorm=False,
residual_in_fp32=False,
zero_centered_weight=False,
is_rms_norm=False,
return_dropout_mask=False,
out_dtype=None,
out=None,
residual_out=None,
):
x_shape_og = x.shape
# reshape input data into 2D tensor
x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1]))
if residual is not None:
assert residual.shape == x_shape_og
residual = maybe_contiguous_lastdim(
residual.reshape(-1, residual.shape[-1])
)
if x1 is not None:
assert x1.shape == x_shape_og
assert rowscale is None, "rowscale is not supported with parallel LayerNorm"
x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1]))
# weight can be None when elementwise_affine=False for LayerNorm
if weight is not None:
weight = weight.contiguous()
bias = maybe_contiguous(bias)
weight1 = maybe_contiguous(weight1)
bias1 = maybe_contiguous(bias1)
if rowscale is not None:
rowscale = rowscale.reshape(-1).contiguous()
residual_dtype = (
residual.dtype
if residual is not None
else (torch.float32 if residual_in_fp32 else None)
)
if out is not None:
out = out.reshape(-1, out.shape[-1])
if residual_out is not None:
residual_out = residual_out.reshape(-1, residual_out.shape[-1])
y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = (
_layer_norm_fwd(
x,
weight,
bias,
eps,
residual,
x1,
weight1,
bias1,
dropout_p=dropout_p,
rowscale=rowscale,
out_dtype=out_dtype,
residual_dtype=residual_dtype,
zero_centered_weight=zero_centered_weight,
is_rms_norm=is_rms_norm,
return_dropout_mask=return_dropout_mask,
out=out,
residual_out=residual_out,
)
)
y = y.reshape(x_shape_og)
if residual is not None:
residual_out = residual_out.reshape(x_shape_og)
return y, residual_out
return y
eps,
out,
y1,
mean,
rstd,
residual=residual,
x1=x1,
weight1=weight1,
bias1=bias1,
residual_out=residual_out,
rowscale=rowscale,
seeds=seeds,
dropout_mask=dropout_mask,
dropout_mask1=dropout_mask1,
dropout_p=dropout_p,
zero_centered_weight=zero_centered_weight,
is_rms_norm=is_rms_norm,
)
return y1, mean, rstd, seeds, dropout_mask, dropout_mask1
@maybe_wrap_jit_kernel_debug
def layer_norm_fn(
def _norm_forward(
x,
weight,
bias,
@@ -473,7 +465,81 @@ def layer_norm_fn(
out=None,
residual_out=None,
):
return LayerNormFn.forward(
x_shape_og = x.shape
# reshape input data into 2D tensor
x = maybe_contiguous_lastdim(x.reshape(-1, x.shape[-1]))
if residual is not None:
assert residual.shape == x_shape_og
residual = maybe_contiguous_lastdim(residual.reshape(-1, residual.shape[-1]))
if x1 is not None:
assert x1.shape == x_shape_og
assert rowscale is None, "rowscale is not supported with parallel LayerNorm"
x1 = maybe_contiguous_lastdim(x1.reshape(-1, x1.shape[-1]))
# weight can be None when elementwise_affine=False for LayerNorm
if weight is not None:
weight = weight.contiguous()
bias = maybe_contiguous(bias)
weight1 = maybe_contiguous(weight1)
bias1 = maybe_contiguous(bias1)
if rowscale is not None:
rowscale = rowscale.reshape(-1).contiguous()
residual_dtype = (
residual.dtype
if residual is not None
else (torch.float32 if residual_in_fp32 else None)
)
if out is not None:
out = out.reshape(-1, out.shape[-1])
if residual_out is not None:
residual_out = residual_out.reshape(-1, residual_out.shape[-1])
y, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1 = (
_layer_norm_fwd(
x,
weight,
bias,
eps,
residual,
x1,
weight1,
bias1,
dropout_p=dropout_p,
rowscale=rowscale,
out_dtype=out_dtype,
residual_dtype=residual_dtype,
zero_centered_weight=zero_centered_weight,
is_rms_norm=is_rms_norm,
return_dropout_mask=return_dropout_mask,
out=out,
residual_out=residual_out,
)
)
y = y.reshape(x_shape_og)
if residual is not None:
residual_out = residual_out.reshape(x_shape_og)
return y, residual_out
return y
def rms_norm_fn(
x,
weight,
bias,
residual=None,
x1=None,
weight1=None,
bias1=None,
eps=1e-6,
dropout_p=0.0,
rowscale=None,
prenorm=False,
residual_in_fp32=False,
zero_centered_weight=False,
return_dropout_mask=False,
out_dtype=None,
out=None,
residual_out=None,
):
return _norm_forward(
x,
weight,
bias,
@@ -487,7 +553,7 @@ def layer_norm_fn(
prenorm,
residual_in_fp32,
zero_centered_weight,
is_rms_norm,
True,
return_dropout_mask,
out_dtype,
out,
@@ -540,7 +606,6 @@ def _norm_infer_kernel(
tl.store(Y + cols, y, mask=cols < N)
@maybe_wrap_jit_kernel_debug
def norm_infer(
x: Tensor,
weight: Optional[Tensor],
@@ -583,100 +648,8 @@ def norm_infer(
return out
@maybe_wrap_jit_kernel_debug
def rms_norm_fn(
x,
weight,
bias,
residual=None,
x1=None,
weight1=None,
bias1=None,
eps=1e-6,
dropout_p=0.0,
rowscale=None,
prenorm=False,
residual_in_fp32=False,
zero_centered_weight=False,
return_dropout_mask=False,
out_dtype=None,
out=None,
residual_out=None,
):
return LayerNormFn.forward(
x,
weight,
bias,
residual,
x1,
weight1,
bias1,
eps,
dropout_p,
rowscale,
prenorm,
residual_in_fp32,
zero_centered_weight,
True,
return_dropout_mask,
out_dtype,
out,
residual_out,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
if current_platform.is_mps():
from .mps_fallback import norm_infer_native, rms_norm_fn_native
@maybe_wrap_jit_kernel_debug
def norm_infer(
x: Tensor,
weight: Optional[Tensor],
bias: Optional[Tensor],
eps: float,
is_rms_norm: bool = False,
out: Optional[Tensor] = None,
):
return norm_infer_native(x, weight, bias, eps, is_rms_norm, out)
@maybe_wrap_jit_kernel_debug
def rms_norm_fn(
x,
weight,
bias,
residual=None,
x1=None,
weight1=None,
bias1=None,
eps=1e-6,
dropout_p=0.0,
rowscale=None,
prenorm=False,
residual_in_fp32=False,
zero_centered_weight=False,
return_dropout_mask=False,
out_dtype=None,
out=None,
residual_out=None,
):
return rms_norm_fn_native(
x,
weight,
bias,
residual,
x1,
weight1,
bias1,
eps,
dropout_p,
rowscale,
prenorm,
residual_in_fp32,
zero_centered_weight,
return_dropout_mask,
out_dtype,
out,
residual_out,
)
norm_infer = norm_infer_native
rms_norm_fn = rms_norm_fn_native

3
python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py
View File

@@ -3,6 +3,7 @@ import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.srt.utils.custom_op import register_custom_op
@@ -69,8 +70,6 @@ def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
return _triton_one_pass_rms_norm_cuda(x, w, eps)
from sglang.multimodal_gen.runtime.platforms import current_platform
if current_platform.is_mps():
from .mps_fallback import triton_one_pass_rms_norm_native

25
python/sglang/jit_kernel/diffusion/triton/rotary.py
View File

@@ -2,7 +2,6 @@ import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -13,7 +12,7 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
triton.Config({"BLOCK_HS_HALF": 128}, num_warps=4),
triton.Config({"BLOCK_HS_HALF": 256}, num_warps=8),
],
key=["head_size", "interleaved"],
key=["head_size"],
)
@triton.jit
def _rotary_embedding_kernel(
@@ -27,7 +26,6 @@ def _rotary_embedding_kernel(
stride_x_row,
stride_cos_row,
stride_sin_row,
interleaved: tl.constexpr,
BLOCK_HS_HALF: tl.constexpr,
):
row_idx = tl.program_id(0)
@@ -65,7 +63,6 @@ def _rotary_embedding_kernel(
tl.store(output_row_ptr + offsets_x2, o2_vals.to(x2_vals.dtype), mask=mask)
@maybe_wrap_jit_kernel_debug
def apply_rotary_embedding(
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
) -> torch.Tensor:
@@ -103,7 +100,6 @@ def apply_rotary_embedding(
x_reshaped.stride(0),
cos.stride(0),
sin.stride(0),
interleaved,
)
return output
@@ -112,24 +108,9 @@ def apply_rotary_embedding(
if current_platform.is_npu():
from .npu_fallback import apply_rotary_embedding_native
@maybe_wrap_jit_kernel_debug
def apply_rotary_embedding(
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
interleaved: bool = False,
) -> torch.Tensor:
return apply_rotary_embedding_native(x, cos, sin, interleaved)
apply_rotary_embedding = apply_rotary_embedding_native
if current_platform.is_mps():
from .mps_fallback import apply_rotary_embedding_native
@maybe_wrap_jit_kernel_debug
def apply_rotary_embedding(
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
interleaved: bool = False,
) -> torch.Tensor:
return apply_rotary_embedding_native(x, cos, sin, interleaved)
apply_rotary_embedding = apply_rotary_embedding_native

228
python/sglang/jit_kernel/diffusion/triton/scale_shift.py
View File

@@ -2,7 +2,6 @@ import torch
import triton # type: ignore
import triton.language as tl # type: ignore
from sglang.jit_kernel.debug_utils import maybe_wrap_jit_kernel_debug
from sglang.multimodal_gen.runtime.platforms import current_platform
@@ -357,95 +356,6 @@ def fuse_scale_shift_kernel_blc_opt(
tl.store(y_ptr + x_off, y, mask=mask)
@triton.jit
def fuse_scale_shift_gate_select01_kernel_blc_opt(
x_ptr,
shift0_ptr,
scale0_ptr,
gate0_ptr,
shift1_ptr,
scale1_ptr,
gate1_ptr,
index_ptr,
y_ptr,
gate_out_ptr,
B,
L,
C,
stride_x_b,
stride_x_l,
stride_x_c,
stride_s0_b,
stride_s0_c,
stride_sc0_b,
stride_sc0_c,
stride_g0_b,
stride_g0_c,
stride_s1_b,
stride_s1_c,
stride_sc1_b,
stride_sc1_c,
stride_g1_b,
stride_g1_c,
stride_i_b,
stride_i_l,
stride_go_b,
stride_go_l,
stride_go_c,
BLOCK_L: tl.constexpr,
BLOCK_C: tl.constexpr,
):
pid_l = tl.program_id(0)
pid_c = tl.program_id(1)
pid_b = tl.program_id(2)
l_offsets = pid_l * BLOCK_L + tl.arange(0, BLOCK_L)
c_offsets = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
mask_l = l_offsets < L
mask_c = c_offsets < C
mask = mask_l[:, None] & mask_c[None, :]
x_off = (
pid_b * stride_x_b
+ l_offsets[:, None] * stride_x_l
+ c_offsets[None, :] * stride_x_c
)
x = tl.load(x_ptr + x_off, mask=mask, other=0)
idx_off = pid_b * stride_i_b + l_offsets * stride_i_l
idx = tl.load(index_ptr + idx_off, mask=mask_l, other=0).to(tl.int1)[:, None]
s0_off = pid_b * stride_s0_b + c_offsets[None, :] * stride_s0_c
sc0_off = pid_b * stride_sc0_b + c_offsets[None, :] * stride_sc0_c
g0_off = pid_b * stride_g0_b + c_offsets[None, :] * stride_g0_c
s1_off = pid_b * stride_s1_b + c_offsets[None, :] * stride_s1_c
sc1_off = pid_b * stride_sc1_b + c_offsets[None, :] * stride_sc1_c
g1_off = pid_b * stride_g1_b + c_offsets[None, :] * stride_g1_c
shift0 = tl.load(shift0_ptr + s0_off, mask=mask_c[None, :], other=0)
scale0 = tl.load(scale0_ptr + sc0_off, mask=mask_c[None, :], other=0)
gate0 = tl.load(gate0_ptr + g0_off, mask=mask_c[None, :], other=0)
shift1 = tl.load(shift1_ptr + s1_off, mask=mask_c[None, :], other=0)
scale1 = tl.load(scale1_ptr + sc1_off, mask=mask_c[None, :], other=0)
gate1 = tl.load(gate1_ptr + g1_off, mask=mask_c[None, :], other=0)
shift = tl.where(idx, shift1, shift0)
scale = tl.where(idx, scale1, scale0)
gate = tl.where(idx, gate1, gate0)
y = x * (1 + scale) + shift
tl.store(y_ptr + x_off, y, mask=mask)
go_off = (
pid_b * stride_go_b
+ l_offsets[:, None] * stride_go_l
+ c_offsets[None, :] * stride_go_c
)
tl.store(gate_out_ptr + go_off, gate, mask=mask)
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_kernel(
x: torch.Tensor,
scale: torch.Tensor,
@@ -465,7 +375,10 @@ def fuse_scale_shift_kernel(
rows = B * L
x_2d = x.view(rows, C)
output_2d = output.view(rows, C)
grid = lambda META: (rows, triton.cdiv(C, META["BLOCK_N"]))
def grid(meta):
return (rows, triton.cdiv(C, meta["BLOCK_N"]))
num_frames = scale.shape[1]
assert (
L % num_frames == 0
@@ -565,80 +478,6 @@ def fuse_scale_shift_kernel(
return output
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_gate_select01_kernel(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
block_l: int = 128,
block_c: int = 128,
):
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]")
grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B)
fuse_scale_shift_gate_select01_kernel_blc_opt[grid](
x,
shift0,
scale0,
gate0,
shift1,
scale1,
gate1,
index,
output,
gate_out,
B,
L,
C,
x.stride(0),
x.stride(1),
x.stride(2),
shift0.stride(0),
shift0.stride(1),
scale0.stride(0),
scale0.stride(1),
gate0.stride(0),
gate0.stride(1),
shift1.stride(0),
shift1.stride(1),
scale1.stride(0),
scale1.stride(1),
gate1.stride(0),
gate1.stride(1),
index.stride(0),
index.stride(1),
gate_out.stride(0),
gate_out.stride(1),
gate_out.stride(2),
BLOCK_L=block_l,
BLOCK_C=block_c,
num_warps=4,
num_stages=2,
)
return output, gate_out
@maybe_wrap_jit_kernel_debug
def fuse_layernorm_scale_shift_gate_select01_kernel(
x: torch.Tensor,
weight: torch.Tensor | None,
@@ -728,7 +567,6 @@ def fuse_layernorm_scale_shift_gate_select01_kernel(
return output, gate_out
@maybe_wrap_jit_kernel_debug
def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
x: torch.Tensor,
residual: torch.Tensor,
@@ -839,61 +677,9 @@ def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
if current_platform.is_npu():
from .npu_fallback import fuse_scale_shift_native
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_kernel(
x: torch.Tensor,
scale: torch.Tensor,
shift: torch.Tensor,
scale_constant: float = 1.0,
block_l: int = 128,
block_c: int = 128,
):
return fuse_scale_shift_native(
x, scale, shift, scale_constant, block_l, block_c
)
fuse_scale_shift_kernel = fuse_scale_shift_native
if current_platform.is_mps():
from .mps_fallback import (
fuse_scale_shift_gate_select01_kernel_native,
fuse_scale_shift_kernel_native,
)
from .mps_fallback import fuse_scale_shift_kernel_native
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_kernel(
x: torch.Tensor,
scale: torch.Tensor,
shift: torch.Tensor,
scale_constant: float = 1.0,
block_l: int = 128,
block_c: int = 128,
):
return fuse_scale_shift_kernel_native(
x, scale, shift, scale_constant, block_l, block_c
)
@maybe_wrap_jit_kernel_debug
def fuse_scale_shift_gate_select01_kernel(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
block_l: int = 128,
block_c: int = 128,
):
return fuse_scale_shift_gate_select01_kernel_native(
x,
scale0,
shift0,
gate0,
scale1,
shift1,
gate1,
index,
block_l,
block_c,
)
fuse_scale_shift_kernel = fuse_scale_shift_kernel_native

41
python/sglang/jit_kernel/tests/test_qwen_image_modulation.py
View File

@@ -6,7 +6,6 @@ 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
@@ -56,17 +55,9 @@ def _baseline_select01_modulation(
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 _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
return output, gate_out
def _baseline_residual_select01_modulation(
@@ -92,19 +83,29 @@ def _baseline_residual_select01_modulation(
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,
output, gate_out = _apply_select01_modulation(
normalized, scale0, shift0, gate0, scale1, shift1, gate1, index
)
return output, residual_out, gate_out
def _apply_select01_modulation(
x: torch.Tensor,
scale0: torch.Tensor,
shift0: torch.Tensor,
gate0: torch.Tensor,
scale1: torch.Tensor,
shift1: torch.Tensor,
gate1: torch.Tensor,
index: torch.Tensor,
):
idx = index.bool().unsqueeze(-1)
scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1))
shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1))
gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1))
return x * (1 + scale) + shift, gate
@pytest.fixture(autouse=True)
def cuda_setup():
if not torch.cuda.is_available():

35
python/sglang/multimodal_gen/runtime/layers/layernorm.py
View File

@@ -10,7 +10,18 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
get_tp_group,
)
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
_is_cuda = current_platform.is_cuda()
_is_npu = current_platform.is_npu()
@@ -24,18 +35,6 @@ if _is_npu:
if _is_musa:
from sgl_kernel import fused_add_rmsnorm
from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn
from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
get_tp_group,
)
from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
# Copied and adapted from sglang
@CustomOp.register("rms_norm")
@@ -68,18 +67,12 @@ class RMSNorm(CustomOp):
x, self.weight, bias=None, residual=residual, eps=self.variance_epsilon
)
def _forward_cuda_fp32_rmsnorm(self, x: torch.Tensor) -> torch.Tensor:
# Avoid wrap_triton in torch.compile: it specializes on a fresh
# constant_args_idx every call and eventually falls back to eager.
return self.forward_native(x)
def forward_cuda(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
shape = x.shape
device = x.device
x = x.reshape(-1, shape[-1])
if residual is not None:
residual_shape = residual.shape
@@ -87,7 +80,7 @@ class RMSNorm(CustomOp):
if x.dtype == torch.float:
if residual is None and self.variance_size_override is None:
return self._forward_cuda_fp32_rmsnorm(x).view(shape)
return self.forward_native(x).view(shape)
out = self.forward_triton(x, residual)
if residual is not None:
return out[0].view(shape), out[1].view(residual_shape)
@@ -215,9 +208,7 @@ class RMSNorm(CustomOp):
return out
def extra_repr(self) -> str:
s = f"hidden_size={self.weight.data.size(0)}"
s += f", eps={self.variance_epsilon}"
return s
return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}"
# Copied and adapted from sglang

6
python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py
View File

@@ -703,11 +703,6 @@ class QwenImageTransformerBlock(nn.Module):
self.quant_config = quant_config
self.zero_cond_t = zero_cond_t
mod_quant_config = (
quant_config
if (quant_config is not None and quant_config.get_name() == "svdquant")
else None
)
# Image processing modules
self.img_mod = nn.Sequential(
nn.SiLU(),
@@ -1102,7 +1097,6 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
@functools.lru_cache(maxsize=50)
def build_modulate_index(self, img_shapes: tuple[int, int, int], device):
sp_world_size = get_sp_world_size()
modulate_index_list = []