From c1fe5de69cfccfb0200eaee478d4ca361b20d75b Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sun, 22 Mar 2026 22:38:57 +0800 Subject: [PATCH] [Diffusion] Clean up diffusion Triton kernels and modernize custom op registration (#21122) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../benchmark/bench_qwen_image_modulation.py | 42 +- .../diffusion/triton/mps_fallback.py | 21 - .../jit_kernel/diffusion/triton/norm.py | 407 ++++++++---------- .../diffusion/triton/rmsnorm_onepass.py | 3 +- .../jit_kernel/diffusion/triton/rotary.py | 25 +- .../diffusion/triton/scale_shift.py | 228 +--------- .../tests/test_qwen_image_modulation.py | 41 +- .../runtime/layers/layernorm.py | 35 +- .../runtime/models/dits/qwen_image.py | 6 - 9 files changed, 257 insertions(+), 551 deletions(-) diff --git a/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py b/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py index e0bc7dee0..42552a805 100644 --- a/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py +++ b/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py @@ -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: diff --git a/python/sglang/jit_kernel/diffusion/triton/mps_fallback.py b/python/sglang/jit_kernel/diffusion/triton/mps_fallback.py index 4ed5b9b36..e1da78bae 100644 --- a/python/sglang/jit_kernel/diffusion/triton/mps_fallback.py +++ b/python/sglang/jit_kernel/diffusion/triton/mps_fallback.py @@ -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: diff --git a/python/sglang/jit_kernel/diffusion/triton/norm.py b/python/sglang/jit_kernel/diffusion/triton/norm.py index 2c717909b..162a87ef7 100644 --- a/python/sglang/jit_kernel/diffusion/triton/norm.py +++ b/python/sglang/jit_kernel/diffusion/triton/norm.py @@ -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 diff --git a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py index 8f3f84939..8c776d7a9 100644 --- a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py +++ b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py @@ -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 diff --git a/python/sglang/jit_kernel/diffusion/triton/rotary.py b/python/sglang/jit_kernel/diffusion/triton/rotary.py index b8942cd6c..02665fbaf 100644 --- a/python/sglang/jit_kernel/diffusion/triton/rotary.py +++ b/python/sglang/jit_kernel/diffusion/triton/rotary.py @@ -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 diff --git a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py index 8d6761f24..2ff71b5d2 100644 --- a/python/sglang/jit_kernel/diffusion/triton/scale_shift.py +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -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 diff --git a/python/sglang/jit_kernel/tests/test_qwen_image_modulation.py b/python/sglang/jit_kernel/tests/test_qwen_image_modulation.py index ae8f7ef40..20ed90078 100644 --- a/python/sglang/jit_kernel/tests/test_qwen_image_modulation.py +++ b/python/sglang/jit_kernel/tests/test_qwen_image_modulation.py @@ -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(): diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 175f48201..62047224f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -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 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 91f963042..ea363021a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -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 = []