diff --git a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py b/python/sglang/jit_kernel/diffusion/triton/norm.py similarity index 52% rename from python/sglang/multimodal_gen/runtime/layers/triton_ops.py rename to python/sglang/jit_kernel/diffusion/triton/norm.py index ddcd78229..17a5bb1ca 100644 --- a/python/sglang/multimodal_gen/runtime/layers/triton_ops.py +++ b/python/sglang/jit_kernel/diffusion/triton/norm.py @@ -1,6 +1,3 @@ -# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo - -# TODO: for temporary usage, expecting a refactor from typing import Optional import torch @@ -8,507 +5,6 @@ import triton # type: ignore import triton.language as tl # type: ignore from torch import Tensor -from sglang.multimodal_gen.runtime.platforms import current_platform - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_N": 64}, num_warps=2), - triton.Config({"BLOCK_N": 128}, num_warps=4), - triton.Config({"BLOCK_N": 256}, num_warps=4), - triton.Config({"BLOCK_N": 512}, num_warps=4), - triton.Config({"BLOCK_N": 1024}, num_warps=8), - ], - key=["inner_dim"], -) -@triton.jit -def _fused_scale_shift_4d_kernel( - output_ptr, - normalized_ptr, - scale_ptr, - shift_ptr, - scale_constant: tl.constexpr, # scale_constant is either 0 or 1. - rows, - inner_dim, - seq_len, - num_frames, - frame_seqlen, - BLOCK_N: tl.constexpr, -): - pid_row = tl.program_id(0) - pid_col = tl.program_id(1) - - col_offsets = pid_col * BLOCK_N + tl.arange(0, BLOCK_N) - mask = col_offsets < inner_dim - - # Pointers for normalized and output - row_base = pid_row * inner_dim - norm_ptrs = normalized_ptr + row_base + col_offsets - out_ptrs = output_ptr + row_base + col_offsets - - # Pointers for scale and shift for 4D - b_idx = pid_row // seq_len - t_idx = pid_row % seq_len - frame_idx_in_batch = t_idx // frame_seqlen - - scale_row_idx = b_idx * num_frames + frame_idx_in_batch - scale_ptrs = scale_ptr + scale_row_idx * inner_dim + col_offsets - shift_ptrs = shift_ptr + scale_row_idx * inner_dim + col_offsets - - normalized = tl.load(norm_ptrs, mask=mask, other=0.0) - scale = tl.load(scale_ptrs, mask=mask, other=0.0) - shift = tl.load(shift_ptrs, mask=mask, other=0.0) - - scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype) - output = normalized * (scale_const_tensor + scale) + shift - - tl.store(out_ptrs, output, mask=mask) - - -@triton.jit -def fuse_scale_shift_kernel_blc_opt( - x_ptr, - shift_ptr, - scale_ptr, - scale_constant: tl.constexpr, # scale_constant is either 0 or 1., - y_ptr, - B, - L, - C, - stride_x_b, - stride_x_l, - stride_x_c, - stride_s_b, - stride_s_l, - stride_s_c, - stride_sc_b, - stride_sc_l, - stride_sc_c, - SCALE_IS_SCALAR: tl.constexpr, - SHIFT_IS_SCALAR: tl.constexpr, - 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) - - if SHIFT_IS_SCALAR: - shift_val = tl.load(shift_ptr) - shift = tl.full((BLOCK_L, BLOCK_C), shift_val, dtype=shift_val.dtype) - else: - s_off = ( - pid_b * stride_s_b - + l_offsets[:, None] * stride_s_l - + c_offsets[None, :] * stride_s_c - ) - shift = tl.load(shift_ptr + s_off, mask=mask, other=0) - - if SCALE_IS_SCALAR: - scale_val = tl.load(scale_ptr) - scale = tl.full((BLOCK_L, BLOCK_C), scale_val, dtype=scale_val.dtype) - else: - sc_off = ( - pid_b * stride_sc_b - + l_offsets[:, None] * stride_sc_l - + c_offsets[None, :] * stride_sc_c - ) - scale = tl.load(scale_ptr + sc_off, mask=mask, other=0) - - y = x * (scale_constant + scale) + shift - 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) - - -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, -): - assert x.is_cuda and scale.is_cuda - assert x.is_contiguous() - - B, L, C = x.shape - output = torch.empty_like(x) - - if scale.dim() == 4: - # scale/shift: [B, F, 1, C] - 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"])) - num_frames = scale.shape[1] - assert ( - L % num_frames == 0 - ), "seq_len must be divisible by num_frames for 4D scale/shift" - frame_seqlen = L // num_frames - - # Compact [B, F, C] without the singleton dim into [B*F, C] - scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous() - shift_reshaped = shift.squeeze(2).reshape(-1, C).contiguous() - - _fused_scale_shift_4d_kernel[grid]( - output_2d, - x_2d, - scale_reshaped, - shift_reshaped, - scale_constant, - rows, - C, - L, - num_frames, - frame_seqlen, - ) - else: - # 2D: [B, C] or [1, C] -> treat as [B, 1, C] and broadcast over L - # 3D: [B, L, C] (or broadcastable variants like [B, 1, C], [1, L, C], [1, 1, C]) - # Also support scalar (0D or 1-element) - if scale.dim() == 0 or (scale.dim() == 1 and scale.numel() == 1): - scale_blc = scale.reshape(1) - elif scale.dim() == 2: - scale_blc = scale[:, None, :] - elif scale.dim() == 3: - scale_blc = scale - else: - raise ValueError("scale must be 0D/1D(1)/2D/3D or 4D") - - if shift.dim() == 0 or (shift.dim() == 1 and shift.numel() == 1): - shift_blc = shift.reshape(1) - elif shift.dim() == 2: - shift_blc = shift[:, None, :] - elif shift.dim() == 3: - shift_blc = shift - else: - # broadcast later via expand if possible - shift_blc = shift - - need_scale_scalar = scale_blc.dim() == 1 and scale_blc.numel() == 1 - need_shift_scalar = shift_blc.dim() == 1 and shift_blc.numel() == 1 - - if not need_scale_scalar: - scale_exp = scale_blc.expand(B, L, C) - s_sb, s_sl, s_sc = scale_exp.stride() - else: - s_sb = s_sl = s_sc = 0 - - if not need_shift_scalar: - shift_exp = shift_blc.expand(B, L, C) - sh_sb, sh_sl, sh_sc = shift_exp.stride() - else: - sh_sb = sh_sl = sh_sc = 0 - - # If both scalars and both zero, copy fast-path - if need_scale_scalar and need_shift_scalar: - if (scale_blc.abs().max() == 0) and (shift_blc.abs().max() == 0): - output.copy_(x) - return output - - grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B) - fuse_scale_shift_kernel_blc_opt[grid]( - x, - shift_blc if need_shift_scalar else shift_exp, - scale_blc if need_scale_scalar else scale_exp, - scale_constant, - output, - B, - L, - C, - x.stride(0), - x.stride(1), - x.stride(2), - sh_sb, - sh_sl, - sh_sc, - s_sb, - s_sl, - s_sc, - SCALE_IS_SCALAR=need_scale_scalar, - SHIFT_IS_SCALAR=need_shift_scalar, - BLOCK_L=block_l, - BLOCK_C=block_c, - num_warps=4, - num_stages=2, - ) - return output - - -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 - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_HS_HALF": 32}, num_warps=2), - triton.Config({"BLOCK_HS_HALF": 64}, num_warps=4), - triton.Config({"BLOCK_HS_HALF": 128}, num_warps=4), - triton.Config({"BLOCK_HS_HALF": 256}, num_warps=8), - ], - key=["head_size", "interleaved"], -) -@triton.jit -def _rotary_embedding_kernel( - output_ptr, - x_ptr, - cos_ptr, - sin_ptr, - num_heads, - head_size, - num_tokens, - stride_x_row, - stride_cos_row, - stride_sin_row, - interleaved: tl.constexpr, - BLOCK_HS_HALF: tl.constexpr, -): - row_idx = tl.program_id(0) - token_idx = (row_idx // num_heads) % num_tokens - - x_row_ptr = x_ptr + row_idx * stride_x_row - cos_row_ptr = cos_ptr + token_idx * stride_cos_row - sin_row_ptr = sin_ptr + token_idx * stride_sin_row - output_row_ptr = output_ptr + row_idx * stride_x_row - - # half size for x1 and x2 - head_size_half = head_size // 2 - - for block_start in range(0, head_size_half, BLOCK_HS_HALF): - offsets_half = block_start + tl.arange(0, BLOCK_HS_HALF) - mask = offsets_half < head_size_half - - cos_vals = tl.load(cos_row_ptr + offsets_half, mask=mask, other=0.0) - sin_vals = tl.load(sin_row_ptr + offsets_half, mask=mask, other=0.0) - - offsets_x1 = 2 * offsets_half - offsets_x2 = 2 * offsets_half + 1 - - x1_vals = tl.load(x_row_ptr + offsets_x1, mask=mask, other=0.0) - x2_vals = tl.load(x_row_ptr + offsets_x2, mask=mask, other=0.0) - - x1_fp32 = x1_vals.to(tl.float32) - x2_fp32 = x2_vals.to(tl.float32) - cos_fp32 = cos_vals.to(tl.float32) - sin_fp32 = sin_vals.to(tl.float32) - o1_vals = tl.fma(-x2_fp32, sin_fp32, x1_fp32 * cos_fp32) - o2_vals = tl.fma(x1_fp32, sin_fp32, x2_fp32 * cos_fp32) - - tl.store(output_row_ptr + offsets_x1, o1_vals.to(x1_vals.dtype), mask=mask) - tl.store(output_row_ptr + offsets_x2, o2_vals.to(x2_vals.dtype), mask=mask) - - -def apply_rotary_embedding( - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False -) -> torch.Tensor: - output = torch.empty_like(x) - - if x.dim() > 3: - bsz, num_tokens, num_heads, head_size = x.shape - else: - num_tokens, num_heads, head_size = x.shape - bsz = 1 - - assert head_size % 2 == 0, "head_size must be divisible by 2" - - x_reshaped = x.view(-1, head_size) - output_reshaped = output.view(-1, head_size) - - # num_tokens per head, 1 token per block - grid = (bsz * num_tokens * num_heads,) - - if interleaved and cos.shape[-1] == head_size: - cos = cos[..., ::2].contiguous() - sin = sin[..., ::2].contiguous() - else: - cos = cos.contiguous() - sin = sin.contiguous() - - _rotary_embedding_kernel[grid]( - output_reshaped, - x_reshaped, - cos, - sin, - num_heads, - head_size, - num_tokens, - x_reshaped.stride(0), - cos.stride(0), - sin.stride(0), - interleaved, - ) - - return output - # RMSNorm-fp32 def maybe_contiguous_lastdim(x): @@ -1122,86 +618,3 @@ def rms_norm_fn( out, residual_out, ) - - -# Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956 -@triton.jit -def _rms_norm_tiled_onepass( - y_ptr, - x_ptr, - w_ptr, - SEQ: tl.constexpr, - DIM: tl.constexpr, - EPS: tl.constexpr, - BLOCK_SIZE_SEQ: tl.constexpr, - BLOCK_SIZE_DIM: tl.constexpr, -): - seq_blk_id = tl.program_id(0) - seq_id = seq_blk_id * BLOCK_SIZE_SEQ - - seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None] - s_mask = seq_offset < SEQ - d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :] - d_mask = d_offset < DIM - y_blk = y_ptr + seq_offset * DIM + d_offset - x_blk = x_ptr + seq_offset * DIM + d_offset - mask = s_mask & d_mask - - x = tl.load(x_blk, mask=mask, other=0.0).to(tl.float32) - mean_square = tl.sum(x * x, axis=1, keep_dims=True) / DIM - rstd = tl.math.rsqrt(mean_square + EPS) - w = tl.load(w_ptr + d_offset, mask=d_mask) - tl.store(y_blk, x * rstd * w, mask=mask) - - -def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6): - shape = x.shape - x = x.contiguous() - y = torch.empty_like(x) - x_view = x.reshape(-1, shape[-1]) - y_view = y.reshape(-1, shape[-1]) - S, D = x_view.shape - - BLOCK_SIZE_SEQ = min(16, triton.next_power_of_2(max(1, S // 512))) - grid = (triton.cdiv(S, BLOCK_SIZE_SEQ),) - - with torch.cuda.device(x.device): - torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid]( - y_view, - x_view, - w, - S, - D, - eps, - BLOCK_SIZE_DIM=triton.next_power_of_2(D), - BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ, - ) - return y - - -if current_platform.is_npu(): - # TODO: remove this when triton ascend bug is fixed - def fuse_scale_shift_native( - x: torch.Tensor, - scale: torch.Tensor, - shift: torch.Tensor, - block_l: int = 128, - block_c: int = 128, - ): - return x * (1 + scale) + shift - - fuse_scale_shift_kernel = fuse_scale_shift_native - - # TODO: remove this when triton ascend bug is fixed - def apply_rotary_embedding_native( - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False - ) -> torch.Tensor: - cos = cos.unsqueeze(-2).to(x.dtype) - sin = sin.unsqueeze(-2).to(x.dtype) - x1 = x[..., ::2] - x2 = x[..., 1::2] - o1 = x1 * cos - x2 * sin - o2 = x2 * cos + x1 * sin - return torch.stack((o1, o2), dim=-1).flatten(-2) - - apply_rotary_embedding = apply_rotary_embedding_native diff --git a/python/sglang/jit_kernel/diffusion/triton/npu_fallback.py b/python/sglang/jit_kernel/diffusion/triton/npu_fallback.py new file mode 100644 index 000000000..9cb65a9d0 --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/npu_fallback.py @@ -0,0 +1,25 @@ +import torch + + +# TODO: remove this when triton ascend bug is fixed +def fuse_scale_shift_native( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + block_l: int = 128, + block_c: int = 128, +): + return x * (1 + scale) + shift + + +# TODO: remove this when triton ascend bug is fixed +def apply_rotary_embedding_native( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False +) -> torch.Tensor: + cos = cos.unsqueeze(-2).to(x.dtype) + sin = sin.unsqueeze(-2).to(x.dtype) + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + return torch.stack((o1, o2), dim=-1).flatten(-2) diff --git a/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py new file mode 100644 index 000000000..df6f7bf21 --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/rmsnorm_onepass.py @@ -0,0 +1,58 @@ +import torch +import triton # type: ignore +import triton.language as tl # type: ignore + + +# Adapted from https://github.com/ModelTC/LightX2V/blob/main/lightx2v/common/ops/norm/triton_ops.py#L905-L956 +@triton.jit +def _rms_norm_tiled_onepass( + y_ptr, + x_ptr, + w_ptr, + SEQ: tl.constexpr, + DIM: tl.constexpr, + EPS: tl.constexpr, + BLOCK_SIZE_SEQ: tl.constexpr, + BLOCK_SIZE_DIM: tl.constexpr, +): + seq_blk_id = tl.program_id(0) + seq_id = seq_blk_id * BLOCK_SIZE_SEQ + + seq_offset = seq_id + tl.arange(0, BLOCK_SIZE_SEQ)[:, None] + s_mask = seq_offset < SEQ + d_offset = tl.arange(0, BLOCK_SIZE_DIM)[None, :] + d_mask = d_offset < DIM + y_blk = y_ptr + seq_offset * DIM + d_offset + x_blk = x_ptr + seq_offset * DIM + d_offset + mask = s_mask & d_mask + + x = tl.load(x_blk, mask=mask, other=0.0).to(tl.float32) + mean_square = tl.sum(x * x, axis=1, keep_dims=True) / DIM + rstd = tl.math.rsqrt(mean_square + EPS) + w = tl.load(w_ptr + d_offset, mask=d_mask) + tl.store(y_blk, x * rstd * w, mask=mask) + + +def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6): + shape = x.shape + x = x.contiguous() + y = torch.empty_like(x) + x_view = x.reshape(-1, shape[-1]) + y_view = y.reshape(-1, shape[-1]) + S, D = x_view.shape + + BLOCK_SIZE_SEQ = min(16, triton.next_power_of_2(max(1, S // 512))) + grid = (triton.cdiv(S, BLOCK_SIZE_SEQ),) + + with torch.get_device_module().device(x.device): + torch.library.wrap_triton(_rms_norm_tiled_onepass)[grid]( + y_view, + x_view, + w, + S, + D, + eps, + BLOCK_SIZE_DIM=triton.next_power_of_2(D), + BLOCK_SIZE_SEQ=BLOCK_SIZE_SEQ, + ) + return y diff --git a/python/sglang/jit_kernel/diffusion/triton/rotary.py b/python/sglang/jit_kernel/diffusion/triton/rotary.py new file mode 100644 index 000000000..067a6ceb4 --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/rotary.py @@ -0,0 +1,113 @@ +import torch +import triton # type: ignore +import triton.language as tl # type: ignore + +from sglang.multimodal_gen.runtime.platforms import current_platform + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_HS_HALF": 32}, num_warps=2), + triton.Config({"BLOCK_HS_HALF": 64}, num_warps=4), + triton.Config({"BLOCK_HS_HALF": 128}, num_warps=4), + triton.Config({"BLOCK_HS_HALF": 256}, num_warps=8), + ], + key=["head_size", "interleaved"], +) +@triton.jit +def _rotary_embedding_kernel( + output_ptr, + x_ptr, + cos_ptr, + sin_ptr, + num_heads, + head_size, + num_tokens, + stride_x_row, + stride_cos_row, + stride_sin_row, + interleaved: tl.constexpr, + BLOCK_HS_HALF: tl.constexpr, +): + row_idx = tl.program_id(0) + token_idx = (row_idx // num_heads) % num_tokens + + x_row_ptr = x_ptr + row_idx * stride_x_row + cos_row_ptr = cos_ptr + token_idx * stride_cos_row + sin_row_ptr = sin_ptr + token_idx * stride_sin_row + output_row_ptr = output_ptr + row_idx * stride_x_row + + # half size for x1 and x2 + head_size_half = head_size // 2 + + for block_start in range(0, head_size_half, BLOCK_HS_HALF): + offsets_half = block_start + tl.arange(0, BLOCK_HS_HALF) + mask = offsets_half < head_size_half + + cos_vals = tl.load(cos_row_ptr + offsets_half, mask=mask, other=0.0) + sin_vals = tl.load(sin_row_ptr + offsets_half, mask=mask, other=0.0) + + offsets_x1 = 2 * offsets_half + offsets_x2 = 2 * offsets_half + 1 + + x1_vals = tl.load(x_row_ptr + offsets_x1, mask=mask, other=0.0) + x2_vals = tl.load(x_row_ptr + offsets_x2, mask=mask, other=0.0) + + x1_fp32 = x1_vals.to(tl.float32) + x2_fp32 = x2_vals.to(tl.float32) + cos_fp32 = cos_vals.to(tl.float32) + sin_fp32 = sin_vals.to(tl.float32) + o1_vals = tl.fma(-x2_fp32, sin_fp32, x1_fp32 * cos_fp32) + o2_vals = tl.fma(x1_fp32, sin_fp32, x2_fp32 * cos_fp32) + + tl.store(output_row_ptr + offsets_x1, o1_vals.to(x1_vals.dtype), mask=mask) + tl.store(output_row_ptr + offsets_x2, o2_vals.to(x2_vals.dtype), mask=mask) + + +def apply_rotary_embedding( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False +) -> torch.Tensor: + output = torch.empty_like(x) + + if x.dim() > 3: + bsz, num_tokens, num_heads, head_size = x.shape + else: + num_tokens, num_heads, head_size = x.shape + bsz = 1 + + assert head_size % 2 == 0, "head_size must be divisible by 2" + + x_reshaped = x.view(-1, head_size) + output_reshaped = output.view(-1, head_size) + + # num_tokens per head, 1 token per block + grid = (bsz * num_tokens * num_heads,) + + if interleaved and cos.shape[-1] == head_size: + cos = cos[..., ::2].contiguous() + sin = sin[..., ::2].contiguous() + else: + cos = cos.contiguous() + sin = sin.contiguous() + + _rotary_embedding_kernel[grid]( + output_reshaped, + x_reshaped, + cos, + sin, + num_heads, + head_size, + num_tokens, + x_reshaped.stride(0), + cos.stride(0), + sin.stride(0), + interleaved, + ) + + return output + + +if current_platform.is_npu(): + from .npu_fallback import apply_rotary_embedding_native + + 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 new file mode 100644 index 000000000..16ff4461e --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/triton/scale_shift.py @@ -0,0 +1,408 @@ +import torch +import triton # type: ignore +import triton.language as tl # type: ignore + +from sglang.multimodal_gen.runtime.platforms import current_platform + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_N": 64}, num_warps=2), + triton.Config({"BLOCK_N": 128}, num_warps=4), + triton.Config({"BLOCK_N": 256}, num_warps=4), + triton.Config({"BLOCK_N": 512}, num_warps=4), + triton.Config({"BLOCK_N": 1024}, num_warps=8), + ], + key=["inner_dim"], +) +@triton.jit +def _fused_scale_shift_4d_kernel( + output_ptr, + normalized_ptr, + scale_ptr, + shift_ptr, + scale_constant: tl.constexpr, # scale_constant is either 0 or 1. + rows, + inner_dim, + seq_len, + num_frames, + frame_seqlen, + BLOCK_N: tl.constexpr, +): + pid_row = tl.program_id(0) + pid_col = tl.program_id(1) + + col_offsets = pid_col * BLOCK_N + tl.arange(0, BLOCK_N) + mask = col_offsets < inner_dim + + # Pointers for normalized and output + row_base = pid_row * inner_dim + norm_ptrs = normalized_ptr + row_base + col_offsets + out_ptrs = output_ptr + row_base + col_offsets + + # Pointers for scale and shift for 4D + b_idx = pid_row // seq_len + t_idx = pid_row % seq_len + frame_idx_in_batch = t_idx // frame_seqlen + + scale_row_idx = b_idx * num_frames + frame_idx_in_batch + scale_ptrs = scale_ptr + scale_row_idx * inner_dim + col_offsets + shift_ptrs = shift_ptr + scale_row_idx * inner_dim + col_offsets + + normalized = tl.load(norm_ptrs, mask=mask, other=0.0) + scale = tl.load(scale_ptrs, mask=mask, other=0.0) + shift = tl.load(shift_ptrs, mask=mask, other=0.0) + + scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype) + output = normalized * (scale_const_tensor + scale) + shift + + tl.store(out_ptrs, output, mask=mask) + + +@triton.jit +def fuse_scale_shift_kernel_blc_opt( + x_ptr, + shift_ptr, + scale_ptr, + scale_constant: tl.constexpr, # scale_constant is either 0 or 1., + y_ptr, + B, + L, + C, + stride_x_b, + stride_x_l, + stride_x_c, + stride_s_b, + stride_s_l, + stride_s_c, + stride_sc_b, + stride_sc_l, + stride_sc_c, + SCALE_IS_SCALAR: tl.constexpr, + SHIFT_IS_SCALAR: tl.constexpr, + 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) + + if SHIFT_IS_SCALAR: + shift_val = tl.load(shift_ptr) + shift = tl.full((BLOCK_L, BLOCK_C), shift_val, dtype=shift_val.dtype) + else: + s_off = ( + pid_b * stride_s_b + + l_offsets[:, None] * stride_s_l + + c_offsets[None, :] * stride_s_c + ) + shift = tl.load(shift_ptr + s_off, mask=mask, other=0) + + if SCALE_IS_SCALAR: + scale_val = tl.load(scale_ptr) + scale = tl.full((BLOCK_L, BLOCK_C), scale_val, dtype=scale_val.dtype) + else: + sc_off = ( + pid_b * stride_sc_b + + l_offsets[:, None] * stride_sc_l + + c_offsets[None, :] * stride_sc_c + ) + scale = tl.load(scale_ptr + sc_off, mask=mask, other=0) + + y = x * (scale_constant + scale) + shift + 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) + + +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, +): + assert x.is_cuda and scale.is_cuda + assert x.is_contiguous() + + B, L, C = x.shape + output = torch.empty_like(x) + + if scale.dim() == 4: + # scale/shift: [B, F, 1, C] + 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"])) + num_frames = scale.shape[1] + assert ( + L % num_frames == 0 + ), "seq_len must be divisible by num_frames for 4D scale/shift" + frame_seqlen = L // num_frames + + # Compact [B, F, C] without the singleton dim into [B*F, C] + scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous() + shift_reshaped = shift.squeeze(2).reshape(-1, C).contiguous() + + _fused_scale_shift_4d_kernel[grid]( + output_2d, + x_2d, + scale_reshaped, + shift_reshaped, + scale_constant, + rows, + C, + L, + num_frames, + frame_seqlen, + ) + else: + # 2D: [B, C] or [1, C] -> treat as [B, 1, C] and broadcast over L + # 3D: [B, L, C] (or broadcastable variants like [B, 1, C], [1, L, C], [1, 1, C]) + # Also support scalar (0D or 1-element) + if scale.dim() == 0 or (scale.dim() == 1 and scale.numel() == 1): + scale_blc = scale.reshape(1) + elif scale.dim() == 2: + scale_blc = scale[:, None, :] + elif scale.dim() == 3: + scale_blc = scale + else: + raise ValueError("scale must be 0D/1D(1)/2D/3D or 4D") + + if shift.dim() == 0 or (shift.dim() == 1 and shift.numel() == 1): + shift_blc = shift.reshape(1) + elif shift.dim() == 2: + shift_blc = shift[:, None, :] + elif shift.dim() == 3: + shift_blc = shift + else: + # broadcast later via expand if possible + shift_blc = shift + + need_scale_scalar = scale_blc.dim() == 1 and scale_blc.numel() == 1 + need_shift_scalar = shift_blc.dim() == 1 and shift_blc.numel() == 1 + + if not need_scale_scalar: + scale_exp = scale_blc.expand(B, L, C) + s_sb, s_sl, s_sc = scale_exp.stride() + else: + s_sb = s_sl = s_sc = 0 + + if not need_shift_scalar: + shift_exp = shift_blc.expand(B, L, C) + sh_sb, sh_sl, sh_sc = shift_exp.stride() + else: + sh_sb = sh_sl = sh_sc = 0 + + # If both scalars and both zero, copy fast-path + if need_scale_scalar and need_shift_scalar: + if (scale_blc.abs().max() == 0) and (shift_blc.abs().max() == 0): + output.copy_(x) + return output + + grid = (triton.cdiv(L, block_l), triton.cdiv(C, block_c), B) + fuse_scale_shift_kernel_blc_opt[grid]( + x, + shift_blc if need_shift_scalar else shift_exp, + scale_blc if need_scale_scalar else scale_exp, + scale_constant, + output, + B, + L, + C, + x.stride(0), + x.stride(1), + x.stride(2), + sh_sb, + sh_sl, + sh_sc, + s_sb, + s_sl, + s_sc, + SCALE_IS_SCALAR=need_scale_scalar, + SHIFT_IS_SCALAR=need_shift_scalar, + BLOCK_L=block_l, + BLOCK_C=block_c, + num_warps=4, + num_stages=2, + ) + return output + + +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 + + +if current_platform.is_npu(): + from .npu_fallback import fuse_scale_shift_native + + fuse_scale_shift_kernel = fuse_scale_shift_native diff --git a/python/sglang/multimodal_gen/runtime/layers/elementwise.py b/python/sglang/multimodal_gen/runtime/layers/elementwise.py index 16db42b90..a9cafb70e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/elementwise.py +++ b/python/sglang/multimodal_gen/runtime/layers/elementwise.py @@ -1,7 +1,7 @@ import torch +from sglang.jit_kernel.diffusion.triton.scale_shift import fuse_scale_shift_kernel from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp -from sglang.multimodal_gen.runtime.layers.triton_ops import fuse_scale_shift_kernel class MulAdd(CustomOp): diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index f92fce78b..888824f60 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -20,6 +20,9 @@ if _is_cuda: if _is_npu: import torch_npu +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, @@ -27,12 +30,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_tp_group, ) from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp -from sglang.multimodal_gen.runtime.layers.triton_ops import ( - fuse_scale_shift_kernel, - norm_infer, - rms_norm_fn, - triton_one_pass_rms_norm, -) from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py index 2a8a65993..903bec2be 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py @@ -32,9 +32,9 @@ from typing import Any, Optional, Tuple import torch +from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_group from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp -from sglang.multimodal_gen.runtime.layers.triton_ops import apply_rotary_embedding from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) 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 2eb446aac..11eae5cea 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -15,6 +15,9 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps 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, +) from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.attention import USPAttention @@ -41,9 +44,6 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( apply_flashinfer_rope_qk_inplace, ) -from sglang.multimodal_gen.runtime.layers.triton_ops import ( - fuse_scale_shift_gate_select01_kernel, -) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum,