[diffusion] refactor: refactor diffusion triton kernels (#18966)

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

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)

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

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

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

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):

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
 
 

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__)

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,