chore: import upstream snapshot with attribution

python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+@cache_once
+def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "diffusion_causal_conv3d_cat_pad",
+        *args,
+        cuda_files=["diffusion/causal_conv3d_cat_pad.cuh"],
+        cuda_wrappers=[
+            (
+                "causal_conv3d_cat_pad",
+                "sglang_causal_conv3d_cat_pad::"
+                f"CausalConv3dCatPadKernel<{args}>::run",
+            )
+        ],
+    )
+
+
+def _causal_conv3d_cat_pad_fake_impl(
+    x: torch.Tensor,
+    cache_x: torch.Tensor,
+    pad_w_left: int,
+    pad_w_right: int,
+    pad_h_top: int,
+    pad_h_bottom: int,
+    pad_d_left: int,
+    pad_d_right: int,
+) -> torch.Tensor:
+    cache_t = cache_x.shape[2]
+    depth_left = pad_d_left - cache_t
+    return torch.empty(
+        (
+            x.shape[0],
+            x.shape[1],
+            x.shape[2] + cache_t + depth_left + pad_d_right,
+            x.shape[3] + pad_h_top + pad_h_bottom,
+            x.shape[4] + pad_w_left + pad_w_right,
+        ),
+        device=x.device,
+        dtype=x.dtype,
+    )
+
+
+@register_custom_op(
+    op_name="diffusion_causal_conv3d_cat_pad",
+    mutates_args=[],
+    fake_impl=_causal_conv3d_cat_pad_fake_impl,
+)
+def _causal_conv3d_cat_pad_custom_op(
+    x: torch.Tensor,
+    cache_x: torch.Tensor,
+    pad_w_left: int,
+    pad_w_right: int,
+    pad_h_top: int,
+    pad_h_bottom: int,
+    pad_d_left: int,
+    pad_d_right: int,
+) -> torch.Tensor:
+    out = _causal_conv3d_cat_pad_fake_impl(
+        x,
+        cache_x,
+        pad_w_left,
+        pad_w_right,
+        pad_h_top,
+        pad_h_bottom,
+        pad_d_left,
+        pad_d_right,
+    )
+    module = _jit_causal_conv3d_cat_pad_module(x.dtype)
+    module.causal_conv3d_cat_pad(
+        out,
+        x,
+        cache_x,
+        pad_w_left,
+        pad_w_right,
+        pad_h_top,
+        pad_h_bottom,
+        pad_d_left,
+        pad_d_right,
+    )
+    return out
+
+
+def fused_causal_conv3d_cat_pad_cuda(
+    x: torch.Tensor,
+    cache_x: torch.Tensor,
+    padding: list[int] | tuple[int, ...],
+) -> torch.Tensor:
+    if x.dtype not in _SUPPORTED_DTYPES:
+        raise RuntimeError(f"unsupported dtype for causal Conv3D cat/pad: {x.dtype}")
+    if not torch.compiler.is_compiling():
+        if (
+            not x.is_cuda
+            or not cache_x.is_cuda
+            or x.dim() != 5
+            or cache_x.dim() != 5
+            or not x.is_contiguous()
+            or not cache_x.is_contiguous()
+            or not can_use_fused_causal_conv3d_cat_pad_cuda(x, cache_x, padding)
+        ):
+            raise RuntimeError("unsupported input for causal Conv3D cat/pad CUDA")
+    return _causal_conv3d_cat_pad_custom_op(x, cache_x, *padding)
+
+
+def can_use_fused_causal_conv3d_cat_pad_cuda(
+    x: torch.Tensor,
+    cache_x: torch.Tensor,
+    padding: list[int] | tuple[int, ...],
+) -> bool:
+    if x.dtype not in _SUPPORTED_DTYPES:
+        return False
+    pad_w_left, pad_w_right, pad_h_top, pad_h_bottom, pad_d_left, pad_d_right = padding
+    cache_t = cache_x.shape[2]
+    depth_left = pad_d_left - cache_t
+    if depth_left < 0 or pad_d_right != 0:
+        return False
+    out_numel = (
+        x.shape[0]
+        * x.shape[1]
+        * (x.shape[2] + cache_t + depth_left + pad_d_right)
+        * (x.shape[3] + pad_h_top + pad_h_bottom)
+        * (x.shape[4] + pad_w_left + pad_w_right)
+    )
+    elem_size = 4 if x.dtype == torch.float32 else 2
+    vec_elems = 16 // elem_size
+    return out_numel % vec_elems == 0

python/sglang/jit_kernel/diffusion/cutedsl/common/norm_fusion.py

@@ -0,0 +1,201 @@
+from typing import Optional, Tuple, Union
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from einops import rearrange
+
+from sglang.jit_kernel.diffusion.cutedsl.common.reduce import (
+    cta_reduce_sum,
+    warp_reduce_sum,
+)
+
+
+@cute.jit
+def apply_norm_cta(
+    norm_type: cutlass.Constexpr,
+    num_warps: cutlass.Constexpr,
+    tidx: cutlass.Int32,
+    tXrX: cute.Tensor,
+    tWrW: Optional[cute.Tensor],
+    tBrB: Optional[cute.Tensor],
+    D: Union[cutlass.Int32, cutlass.Constexpr],
+    eps: Union[cutlass.Float32, cutlass.Constexpr],
+) -> cute.Tensor:
+    if cutlass.const_expr(norm_type == "rms"):
+        return apply_rmsnorm_cta(num_warps, tidx, tXrX, tWrW, D, eps)
+    else:
+        return apply_layernorm_cta(num_warps, tidx, tXrX, tWrW, tBrB, D, eps)
+
+
+@cute.jit
+def apply_rmsnorm_cta(
+    num_warps: Union[cutlass.Int32, cutlass.Constexpr],
+    tidx: cutlass.Int32,
+    tXrX: cute.Tensor,
+    tWrW: Optional[cute.Tensor],
+    D: Union[cutlass.Int32, cutlass.Constexpr],
+    eps: Union[cutlass.Float32, cutlass.Constexpr],
+) -> cute.Tensor:
+    """
+    RMSNorm:
+      y[i] = x[i] / sqrt(sum(x ^ 2) / D + eps) * w[i]
+    """
+    val = cute.Float32(0.0)
+    for idx in range(cute.size(tXrX)):
+        # Accumulate in FP32 to improve numerical precision.
+        x_fp32 = tXrX[idx].to(cutlass.Float32)
+        val += x_fp32 * x_fp32
+    val = warp_reduce_sum(val)
+    acc_sq = cta_reduce_sum(val, num_warps, tidx)
+    factor = cute.rsqrt(acc_sq / D + eps)
+    tNrN = cute.make_fragment_like(tXrX)
+    if cutlass.const_expr(isinstance(tWrW, cute.Tensor)):
+        tNrN.store((tXrX.load() * factor * tWrW.load()).to(tNrN.element_type))
+    else:
+        tNrN.store((tXrX.load() * factor).to(tNrN.element_type))
+    return tNrN
+
+
+@cute.jit
+def apply_layernorm_cta(
+    num_warps: Union[cutlass.Int32, cutlass.Constexpr],
+    tidx: cutlass.Int32,
+    tXrX: cute.Tensor,
+    tWrW: Optional[cute.Tensor],
+    tBrB: Optional[cute.Tensor],
+    D: Union[cutlass.Int32, cutlass.Constexpr],
+    eps: Union[cutlass.Float32, cutlass.Constexpr],
+) -> cute.Tensor:
+    """
+    LayerNorm:
+        mean = sum(x) / D
+        var  = sum((x - mean) ^ 2) / D
+        y[i] = (x[i] - mean) / sqrt(var + eps) * w[i] + b[i]
+    """
+    # Reduce mean
+    val = cute.Float32(0.0)
+    for idx in range(cute.size(tXrX)):
+        # Accumulate in FP32 to improve numerical precision.
+        val += tXrX[idx].to(cutlass.Float32)
+    val = warp_reduce_sum(val)
+    val = cta_reduce_sum(val, num_warps, tidx)
+    mean = val / D
+    # Reduce variance
+    val = cute.Float32(0.0)
+    for idx in range(cute.size(tXrX)):
+        # Accumulate in FP32 to improve numerical precision.
+        x_fp32 = tXrX[idx].to(cutlass.Float32)
+        val += (x_fp32 - mean) * (x_fp32 - mean)
+    val = warp_reduce_sum(val)
+    val = cta_reduce_sum(val, num_warps, tidx)
+    factor = cute.rsqrt(val / D + eps)
+    # Normalize
+    tNrN = cute.make_fragment_like(tXrX)
+    if cutlass.const_expr(
+        isinstance(tWrW, cute.Tensor) and isinstance(tBrB, cute.Tensor)
+    ):
+        tNrN.store(
+            ((tXrX.load() - mean) * factor * tWrW.load() + tBrB.load()).to(
+                tNrN.element_type
+            )
+        )
+    else:
+        tNrN.store(((tXrX.load() - mean) * factor).to(tNrN.element_type))
+    return tNrN
+
+
+################################################################################
+# BSFD Indexing
+################################################################################
+# In diffusion norm-fusion kernels, we compute `norm(x) + y`, where
+# `x` has shape [B, S, D] and `y` may come in various broadcastable forms:
+#   [1], [D], [1, D], [1, 1, D], [B, D], [B, 1, D], [B, S, D], or [B, F, 1, D].
+#
+# For a given (batch_id, seq_id), the index mapping for `y` falls into 3 cases:
+#   1) Scalar broadcast [1]:
+#        (batch_id, seq_id, *) -> (0)
+#   2) Frame-based BSFD broadcast [B, F, 1, D]:
+#        frame_id = seq_id // len_frame
+#        (batch_id, seq_id, *) -> (batch_id, frame_id, *)
+#   3) All other cases:
+#        `y` is broadcast to [B, S, D] (via view/expand, no materialization),
+#        and indexed as (batch_id, seq_id, *).
+#
+# This helper normalizes `y` into a BSFD-compatible view so that kernel
+# indexing logic remains simple and uniform.
+################################################################################
+
+
+def broadcast_tensor_for_bsfd(
+    tensor: Union[Optional[torch.Tensor], int],
+    B: int,
+    S: int,
+    D: int,
+) -> Union[Optional[torch.Tensor], int]:
+    """
+    Broadcast to (B, S, D) without memory copy for following shapes:
+    - [D], [1, D], [1, 1, D], [B, D], [B, 1, D], [B, S, D].
+    """
+
+    # Return directly for non-tensor value
+    if not isinstance(tensor, torch.Tensor):
+        return tensor
+
+    if tensor.ndim == 1:
+        # Scalar [1] is preserved as-is and handled specially in CuTe kernel.
+        if tensor.numel() == 1:
+            return tensor
+        return rearrange(tensor, "d -> 1 1 d").expand(B, S, D)
+    if tensor.ndim == 2:
+        return rearrange(tensor, "b d -> b 1 d").expand(B, S, D)
+    if tensor.ndim == 3:
+        return tensor.expand(B, S, D)
+    if tensor.ndim == 4:
+        return tensor
+    raise ValueError(f"BSFD broadcast: unsupported tensor ndim: {tensor.ndim}.")
+
+
+@cute.jit
+def tensor_slice_for_bsfd(
+    mV: cute.Tensor,
+    thr_copy: cute.ThrCopy,
+    batch_id: cutlass.Int32,
+    seq_id: cutlass.Int32,
+    S: Union[cutlass.Int32, cutlass.Constexpr],
+    D: Union[cutlass.Int32, cutlass.Constexpr],
+) -> Tuple[cute.Tensor, cute.Tensor]:
+    """
+    Slice a BSFD-compatible tensor into a per-thread gmem tile and rmem fragment.
+
+    Given a logical (batch_id, seq_id), this helper selects the corresponding
+    D-length slice from `mV` and prepares it for vectorized copy.
+    """
+    gV: cute.Tensor
+    if cutlass.const_expr(cute.is_static(mV.layout) and cute.size(mV.layout) == 1):
+        # build a ((1,1),(1,)) layout so it could broadcast-align with the
+        # regular rmem fragment shape ((4,1),(k,)).
+        layout = cute.make_layout(shape=((1, 1), (1,)))
+        tVgV = cute.make_tensor(mV.iterator, layout)
+        tVrV = cute.make_rmem_tensor(layout, mV.element_type)
+        return tVgV, tVrV
+
+    # Use `local_tile` instead of direct indexing to preserve gmem base pointer
+    # alignment required for vectorized loads.
+    if cutlass.const_expr(len(mV.shape) == 1):
+        gV = mV
+    elif cutlass.const_expr(len(mV.shape) == 3):
+        gV = cute.local_tile(mV, tiler=(1, 1, D), coord=(batch_id, seq_id, 0))
+        gV = gV[0, 0, None]
+    elif cutlass.const_expr(len(mV.shape) == 4):
+        # Compute frame length at runtime (instead of compile time) to avoid
+        # specializing kernels on the frame dimension.
+        frame_len = S // mV.shape[1]
+        frame_id = seq_id // frame_len
+        gV = cute.local_tile(mV, tiler=(1, 1, 1, D), coord=(batch_id, frame_id, 0, 0))
+        gV = gV[0, 0, 0, None]
+    else:
+        raise NotImplementedError(f"BSFD slice: unsupported shape {mV.shape}.")
+    tVgV = thr_copy.partition_S(gV)
+    tVrV = cute.make_fragment_like(tVgV, tVgV.element_type)
+    return tVgV, tVrV

python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py

@@ -0,0 +1,33 @@
+import math
+
+import cutlass
+import cutlass.cute as cute
+
+
+@cute.jit
+def warp_reduce_sum(val: cute.Numeric, reduce_size: int = 32) -> cute.Numeric:
+    iters = int(math.log2(reduce_size))
+    for i in range(iters):
+        val = val + cute.arch.shuffle_sync_down(val, offset=1 << (iters - i - 1))
+    return val
+
+
+@cute.jit
+def cta_reduce_sum(
+    val: cute.Numeric, num_warps: cutlass.Constexpr, tidx: cutlass.Int32
+) -> cute.Numeric:
+    smem = cutlass.utils.SmemAllocator()
+    acc = smem.allocate_tensor(cutlass.Float32, num_warps + 1)
+    warp_id = tidx >> 5
+    lane_id = tidx & 31
+    if lane_id == 0:
+        acc[warp_id] = val
+    cute.arch.sync_threads()
+    if warp_id == 0:
+        val = acc[lane_id] if lane_id < num_warps else cutlass.Float32(0)
+        val = warp_reduce_sum(val)
+        if lane_id == 0:
+            acc[num_warps] = val
+    cute.arch.sync_threads()
+    val = acc[num_warps]
+    return val

python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py

@@ -0,0 +1,344 @@
+from typing import Optional, Tuple
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+
+from sglang.jit_kernel.diffusion.cutedsl.common.norm_fusion import (
+    apply_norm_cta,
+    broadcast_tensor_for_bsfd,
+    tensor_slice_for_bsfd,
+)
+from sglang.jit_kernel.diffusion.cutedsl.utils import (
+    WARP_SIZE,
+    to_cute_arg,
+    to_fake_cute_args,
+)
+
+_COMPILE_CACHE = {}
+
+
+class NormTanhMulAddNormScale:
+    @classmethod
+    def make_hash_key(cls, *inputs):
+        """
+        Compile-time values:
+          - D: hidden dimension (size of the last dimension)
+          - norm_type: layer norm or RMS norm
+          - tensor dtype
+          - tensor rank (i.e., tensor.ndim)
+
+        Runtime values:
+          - all other inputs
+
+        This hash key defines the compile-time specialization boundary for
+        NormTanhMulAddNormScale kernels.
+        """
+
+        def _sig(val):
+            if isinstance(val, torch.Tensor):
+                return (val.dtype, val.ndim, val.shape[-1])
+            return val
+
+        return tuple(_sig(val) for val in inputs)
+
+    def __init__(self, D: int, norm_type: str, is_norm2: bool):
+        self.D = D
+        self.norm_type = norm_type  # "layer" or "rms"
+        self.is_norm2 = is_norm2  # single norm or double norm
+        self.num_warps = self.D // 256  # num of warps per cta
+        self.num_threads = self.num_warps * WARP_SIZE  # num of threads per cta
+
+    @cute.jit
+    def __call__(
+        self,
+        mY,
+        mY2,
+        mX,
+        mWeight,
+        mBias,
+        mScale,
+        mShift,
+        mWeight2,
+        mBias2,
+        mScale2,
+        eps: cutlass.Float32 = cutlass.Float32(1e-5),
+        stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
+    ):
+        # Tensor shapes
+        B, S, _ = mX.shape  # (batch, seq_len, hidden_dim)
+        # Vectorized copy configuration
+        num_vectorized = 8  # maximum num of elem per copy
+        atom_copy = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mX.element_type,
+            num_bits_per_copy=128,
+        )
+        # Thread/value layouts for tiled copy
+        t_layout = cute.make_layout(self.num_threads)  # thread layout within a CTA
+        v_layout = cute.make_layout(num_vectorized)  # per-thread vector layout
+        tiled_copy = cute.make_tiled_copy_tv(atom_copy, t_layout, v_layout)
+
+        self.kernel(
+            mY,
+            mY2,
+            mX,
+            mWeight,
+            mBias,
+            mScale,
+            mShift,
+            mWeight2,
+            mBias2,
+            mScale2,
+            tiled_copy,
+            eps,
+        ).launch(
+            grid=[B * S, 1, 1],
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mY,
+        mY2,
+        mX,
+        mWeight,
+        mBias,
+        mScale,
+        mShift,
+        mWeight2,
+        mBias2,
+        mScale2,
+        tiled_copy: cute.TiledCopy,
+        eps: cutlass.Float32,
+    ):
+        _, S, _ = mX.shape
+        tidx, _, _ = cute.arch.thread_idx()  # thread index
+        bid, _, _ = cute.arch.block_idx()  # cta index
+        bidx = cutlass.Int32(bid // S)  # batch index
+        bidy = cutlass.Int32(bid % S)  # seq_len index
+        thr_copy = tiled_copy.get_slice(tidx)
+
+        @cute.jit
+        def slice_if(mV):
+            if cutlass.const_expr(isinstance(mV, cute.Tensor)):
+                return tensor_slice_for_bsfd(mV, thr_copy, bidx, bidy, S, self.D)
+            return mV, mV
+
+        @cute.jit
+        def copy_if(src, dst):
+            if cutlass.const_expr(
+                isinstance(src, cute.Tensor) and isinstance(dst, cute.Tensor)
+            ):
+                cute.autovec_copy(src, dst)  # LDG.128
+
+        @cute.jit
+        def norm(x, weight, bias):
+            return apply_norm_cta(
+                self.norm_type, self.num_warps, tidx, x, weight, bias, self.D, eps
+            )
+
+        # Slice: retrieve the per-thread data slices for both global memory (gmem)
+        tXgX, tXrX = slice_if(mX)  # x
+        tWgW, tWrW = slice_if(mWeight)  # weight
+        tBgB, tBrB = slice_if(mBias)  # bias
+        tSCgSC, tSCrSC = slice_if(mScale)  # scale
+        tSHgSH, tSHrSH = slice_if(mShift)  # shift
+        tYgY, tYrY = slice_if(mY)  # y
+        if cutlass.const_expr(self.is_norm2):
+            tYgY2, tYrY2 = slice_if(mY2)  # y2
+            tWgW2, tWrW2 = slice_if(mWeight2)  # weight2
+            tBgB2, tBrB2 = slice_if(mBias2)  # bias2
+            tSCgSC2, tSCrSC2 = slice_if(mScale2)  # scale2
+        # Load: load tensor from global memory to registers
+        copy_if(tXgX, tXrX)  # gmem -> rmem
+        copy_if(tWgW, tWrW)  # gmem -> rmem
+        copy_if(tBgB, tBrB)  # gmem -> rmem
+        tNrN = norm(tXrX, tWrW, tBrB)
+        # Compute: value = value * tanh(<scale>) + <shift>
+        copy_if(tSCgSC, tSCrSC)  # gmem -> rmem
+        copy_if(tSHgSH, tSHrSH)  # gmem -> rmem
+        value = tNrN.load() * cute.tanh(tSCrSC.load()) + tSHrSH.load()
+        # Store: y
+        tYrY.store(value.to(tYrY.element_type))
+        copy_if(tYrY, tYgY)  # rmem -> gmem
+        if cutlass.const_expr(self.is_norm2):
+            copy_if(tWgW2, tWrW2)  # gmem -> rmem
+            copy_if(tBgB2, tBrB2)  # gmem -> rmem
+            tNrN2 = norm(tYrY, tWrW2, tBrB2)
+            # Compute: value2 = value2 * (1 + <scale2>)
+            copy_if(tSCgSC2, tSCrSC2)  # gmem -> rmem
+            value2 = tNrN2.load() * (1 + tSCrSC2.load())
+            # Store: y2
+            tYrY2.store(value2.to(tYrY2.element_type))
+            copy_if(tYrY2, tYgY2)  # rmem -> gmem
+
+
+def validate_3d(t: torch.Tensor, B: int, S: int, D: int):
+    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
+    if (
+        t.ndim != 3
+        or (t.shape[0] not in (1, B))
+        or (t.shape[1] not in (1, S) or t.shape[2] != D)
+    ):
+        raise ValueError(f"Validate failed: unsupported 3d-tensor: {t.shape}.")
+    if t.stride()[-1] != 1:
+        raise ValueError(f"Validate failed: not contiguous on dim D.")
+
+
+def validate_weight_bias(t: Optional[torch.Tensor], D: int):
+    if t is None:
+        return
+    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
+    if t.shape != (D,):
+        raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
+    if t.stride()[-1] != 1:
+        raise ValueError(f"Validate failed: not contiguous on dim D.")
+
+
+@torch.library.custom_op("sglang::fused_norm_tanh_mul_add", mutates_args=())
+def fused_norm_tanh_mul_add(
+    x: torch.Tensor,
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    norm_type: str,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    """
+    Fuse: norm(x) * tanh(scale) + shift
+      where norm is either layernorm or rmsnorm.
+
+    Expects:
+      - x: [B, S, D]
+      - weight/bias: None, [D]
+      - scale/shift: [1/B, 1/S, D]
+      - norm_type: str, "layer" or "rms"
+      - eps: Optional[float], default: 1e-5
+
+    D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
+    thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
+    """
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    # Tensor Validation
+    BSD = x.shape
+    validate_3d(x, *BSD)
+    validate_weight_bias(weight, BSD[2])
+    validate_weight_bias(bias, BSD[2])
+    validate_3d(scale, *BSD)
+    validate_3d(shift, *BSD)
+    if norm_type == "layer" or norm_type == "rms":
+        D = x.shape[-1]
+        if D % 256 != 0 or D > 8192:
+            raise ValueError(
+                f"D={D} not supported, must be multiple of 256 and <= 8192"
+            )
+        y = torch.empty_like(x)  # create output tensor
+        scale = broadcast_tensor_for_bsfd(scale, *x.shape)  # handle various shapes
+        shift = broadcast_tensor_for_bsfd(shift, *x.shape)  # handle various shapes
+        # y2, weight2, bias2, scale2 is None
+        torch_tensors = [y, None, x, weight, bias, scale, shift, None, None, None]
+        cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
+        # Compile cache
+        hash_key = NormTanhMulAddNormScale.make_hash_key(norm_type, *torch_tensors)
+        compiled_fn = _COMPILE_CACHE.get(hash_key)
+        if compiled_fn is None:
+            kernel = NormTanhMulAddNormScale(D, norm_type, is_norm2=False)
+            fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
+            compiled_fn = cute.compile(
+                kernel, *fake_sig_args, options="--enable-tvm-ffi"
+            )
+            _COMPILE_CACHE[hash_key] = compiled_fn
+        # Execute
+        compiled_fn(*cute_tensor_args, eps, stream)
+        return y
+    else:
+        raise ValueError(f'norm_type must be one of "layer" and "rms"')
+
+
+@fused_norm_tanh_mul_add.register_fake
+def _fused_norm_tanh_mul_add_fake(x, weight, bias, scale, shift, norm_type, eps=1e-5):
+    return x.new_empty(x.shape)
+
+
+@torch.library.custom_op("sglang::fused_norm_tanh_mul_add_norm_scale", mutates_args=())
+def fused_norm_tanh_mul_add_norm_scale(
+    x: torch.Tensor,
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    weight2: Optional[torch.Tensor],
+    bias2: Optional[torch.Tensor],
+    scale2: torch.Tensor,
+    norm_type: str,
+    eps: float = 1e-5,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Fuse:
+      y = norm(x) * tanh(scale) + shift
+      y2 = norm(y) * (1 + scale2)
+      where norm is either layernorm or rmsnorm.
+
+    Expects:
+      - x: [B, S, D]
+      - weight/bia/weight2/bias2: None, [D]
+      - scale/shift/scale2: [1/B, 1/S, D]
+      - norm_type: str, "layer" or "rms"
+      - eps: Optional[float], default: 1e-5
+
+    D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
+    thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
+    """
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    # Tensor Validation
+    BSD = x.shape
+    validate_3d(x, *BSD)
+    validate_weight_bias(weight, BSD[2])
+    validate_weight_bias(bias, BSD[2])
+    validate_3d(scale, *BSD)
+    validate_3d(shift, *BSD)
+    validate_weight_bias(weight2, BSD[2])
+    validate_weight_bias(bias2, BSD[2])
+    validate_3d(scale2, *BSD)
+    if norm_type == "layer" or norm_type == "rms":
+        D = x.shape[-1]
+        if D % 256 != 0 or D > 8192:
+            raise ValueError(
+                f"D={D} not supported, must be multiple of 256 and <= 8192"
+            )
+        y = torch.empty_like(x)  # create output tensor
+        y2 = torch.empty_like(x)  # create output tensor
+        scale = broadcast_tensor_for_bsfd(scale, *x.shape)  # handle various shapes
+        shift = broadcast_tensor_for_bsfd(shift, *x.shape)  # handle various shapes
+        scale2 = broadcast_tensor_for_bsfd(scale2, *x.shape)  # handle various shapes
+        torch_tensors = [y, y2, x, weight, bias, scale, shift, weight2, bias2, scale2]
+        cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
+        # Compile cache
+        hash_key = NormTanhMulAddNormScale.make_hash_key(norm_type, *torch_tensors)
+        compiled_fn = _COMPILE_CACHE.get(hash_key)
+        if compiled_fn is None:
+            kernel = NormTanhMulAddNormScale(D, norm_type, is_norm2=True)
+            fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
+            compiled_fn = cute.compile(
+                kernel, *fake_sig_args, options="--enable-tvm-ffi"
+            )
+            _COMPILE_CACHE[hash_key] = compiled_fn
+        # Execute
+        compiled_fn(*cute_tensor_args, eps, stream)
+        return y, y2
+    else:
+        raise ValueError(f'norm_type must be one of "layer" and "rms"')
+
+
+@fused_norm_tanh_mul_add_norm_scale.register_fake
+def _fused_norm_tanh_mul_add_norm_scale_fake(
+    x, weight, bias, scale, shift, weight2, bias2, scale2, norm_type, eps=1e-5
+):
+    return x.new_empty(x.shape), x.new_empty(x.shape)

python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py

@@ -0,0 +1,413 @@
+from typing import Optional, Tuple, Union
+
+import cuda.bindings.driver as cuda
+import cutlass
+import cutlass.cute as cute
+import torch
+
+from sglang.jit_kernel.diffusion.cutedsl.common.norm_fusion import (
+    apply_norm_cta,
+    broadcast_tensor_for_bsfd,
+    tensor_slice_for_bsfd,
+)
+from sglang.jit_kernel.diffusion.cutedsl.utils import (
+    WARP_SIZE,
+    to_fake_cute_args,
+)
+
+_COMPILE_CACHE = {}
+
+
+class ScaleResidualNormScaleShift:
+    @classmethod
+    def make_hash_key(cls, *inputs):
+        """
+        Compile-time values:
+          - D: hidden dimension (size of the last dimension)
+          - norm_type: layer norm or RMS norm
+          - tensor dtype
+          - tensor rank (i.e., tensor.ndim)
+
+        Runtime values:
+          - all other inputs
+
+        This hash key defines the compile-time specialization boundary for
+        ScaleResidualNormScaleShift kernels.
+        """
+
+        def _sig(val):
+            if isinstance(val, torch.Tensor):
+                return (val.dtype, val.ndim, val.shape[-1])
+            return val
+
+        return tuple(_sig(val) for val in inputs)
+
+    def __init__(self, D: int, norm_type: str):
+        self.D = D
+        self.norm_type = norm_type  # "layer" or "rms"
+        self.num_warps = self.D // 256  # num of warps per cta
+        self.num_threads = self.num_warps * WARP_SIZE  # num of threads per cta
+
+    @cute.jit
+    def __call__(
+        self,
+        mY,
+        mResOut,
+        mRes,
+        mX,
+        mGate,
+        mWeight,
+        mBias,
+        mScale,
+        mShift,
+        eps: cutlass.Float32 = cutlass.Float32(1e-5),
+        stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
+    ):
+        # Tensor shapes
+        B, S, _ = mX.shape  # (batch, seq_len, hidden_dim)
+        # Vectorized copy configuration
+        num_vectorized = 8  # maximum num of elem per copy
+        atom_copy = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            mX.element_type,
+            num_bits_per_copy=128,
+        )
+        # Thread/value layouts for tiled copy
+        t_layout = cute.make_layout(self.num_threads)  # thread layout within a CTA
+        v_layout = cute.make_layout(num_vectorized)  # per-thread vector layout
+        tiled_copy = cute.make_tiled_copy_tv(atom_copy, t_layout, v_layout)
+
+        self.kernel(
+            mY,
+            mResOut,
+            mRes,
+            mX,
+            mGate,
+            mWeight,
+            mBias,
+            mScale,
+            mShift,
+            tiled_copy,
+            eps,
+        ).launch(
+            grid=[B * S, 1, 1],
+            block=[self.num_threads, 1, 1],
+            stream=stream,
+        )
+
+    @cute.kernel
+    def kernel(
+        self,
+        mY,
+        mResOut,
+        mRes,
+        mX,
+        mGate,
+        mWeight,
+        mBias,
+        mScale,
+        mShift,
+        tiled_copy: cute.TiledCopy,
+        eps: cutlass.Float32,
+    ):
+        _, S, _ = mX.shape
+        tidx, _, _ = cute.arch.thread_idx()  # thread index
+        bid, _, _ = cute.arch.block_idx()  # cta index
+        bidx = cutlass.Int32(bid // S)  # batch index
+        bidy = cutlass.Int32(bid % S)  # seq_len index
+        thr_copy = tiled_copy.get_slice(tidx)
+
+        @cute.jit
+        def slice_if(mV):
+            if cutlass.const_expr(isinstance(mV, cute.Tensor)):
+                return tensor_slice_for_bsfd(mV, thr_copy, bidx, bidy, S, self.D)
+            return mV, mV
+
+        @cute.jit
+        def copy_if(src, dst):
+            if cutlass.const_expr(
+                isinstance(src, cute.Tensor) and isinstance(dst, cute.Tensor)
+            ):
+                cute.autovec_copy(src, dst)  # LDG.128
+
+        @cute.jit
+        def norm(x, weight, bias):
+            return apply_norm_cta(
+                self.norm_type, self.num_warps, tidx, x, weight, bias, self.D, eps
+            )
+
+        # Slice: retrieve the per-thread data slices for both global memory (gmem)
+        # and register memory (rmem). The layouts are:
+        # - ((4,2),(1)):((1,4),(0)) for fp32
+        # - ((8,1),(1)):((1,0),(0)) for fp16/bf16
+        tRgR, tRrR = slice_if(mRes)  # residual
+        tXgX, tXrX = slice_if(mX)  # x
+        tGgG, tGrG = slice_if(mGate)  # gate
+        tROgRO, tROrRO = slice_if(mResOut)  # residual_out
+        tWgW, tWrW = slice_if(mWeight)  # weight
+        tBgB, tBrB = slice_if(mBias)  # bias
+        tSCgSC, tSCrSC = slice_if(mScale)  # scale
+        tSHgSH, tSHrSH = slice_if(mShift)  # shift
+        tYgY, tYrY = slice_if(mY)  # y
+        # Load: load tensor from global memory to registers
+        copy_if(tRgR, tRrR)  # gmem -> rmem
+        copy_if(tXgX, tXrX)  # gmem -> rmem
+        copy_if(tGgG, tGrG)  # gmem -> rmem
+        copy_if(tWgW, tWrW)  # gmem -> rmem
+        copy_if(tBgB, tBrB)  # gmem -> rmem
+
+        # For norm_scale_shift, output:
+        # - y = norm(x, weight, bias) * (1 + scale) + shift
+        # For scale_residual_norm_scale_shift, output:
+        # - residual_out = residual + gate * x
+        # - y = norm(residual_out, weight, bias) * (1 + scale) + shift
+        # Compute: value = <gate> * x
+        value = tXrX.load()
+        if cutlass.const_expr(isinstance(tGrG, cute.Tensor)):
+            value = tGrG.load() * value
+        # Compute: value = value + <residual>
+        if cutlass.const_expr(isinstance(tRrR, cute.Tensor)):
+            value = value + tRrR.load()
+        # Store: residual_out
+        if cutlass.const_expr(isinstance(tROrRO, cute.Tensor)):
+            tROrRO.store(value.to(tROrRO.element_type))
+            copy_if(tROrRO, tROgRO)  # rmem -> gmem
+        # Compute: value = norm(value) * <weight> + <bias>
+        tNrN = cute.make_rmem_tensor_like(tXrX, tXrX.element_type)
+        tNrN.store(value.to(tNrN.element_type))
+        tNrN = norm(tNrN, tWrW, tBrB)
+        # Compute: value = value * (1 + <scale>) + <shift>
+        value = tNrN.load()
+        copy_if(tSCgSC, tSCrSC)  # gmem -> rmem
+        copy_if(tSHgSH, tSHrSH)  # gmem -> rmem
+        if cutlass.const_expr(isinstance(tSCrSC, cute.Tensor)):
+            value = value * (1 + tSCrSC.load())
+        if cutlass.const_expr(isinstance(tSHrSH, cute.Tensor)):
+            value = value + tSHrSH.load()
+        # Store: y
+        tYrY.store(value.to(tYrY.element_type))
+        copy_if(tYrY, tYgY)  # rmem -> gmem
+
+
+def validate_x(t: torch.Tensor, B: int, S: int, D: int):
+    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
+    if t.shape != (B, S, D):
+        raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
+    if t.stride()[-1] != 1:
+        raise ValueError(f"Validate failed: not contiguous on dim D.")
+
+
+def validate_weight_bias(t: Optional[torch.Tensor], D: int):
+    if t is None:
+        return
+    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
+    if t.shape != (D,):
+        raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
+    if t.stride()[-1] != 1:
+        raise ValueError(f"Validate failed: not contiguous on dim D.")
+
+
+def validate_scale_shift(t: torch.Tensor, B: int, S: int, D: int):
+    if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
+    failed = False
+    if t.ndim == 1 and (t.shape[0] not in (1, D)):
+        failed = True
+    elif t.ndim == 2 and ((t.shape[0] not in (1, B)) or t.shape[1] != D):
+        failed = True
+    elif t.ndim == 3 and (
+        (t.shape[0] not in (1, B)) or (t.shape[1] not in (1, S) or t.shape[2] != D)
+    ):
+        failed = True
+    elif t.ndim == 4:
+        F = t.shape[1]
+        if t.shape[0] != B or t.shape[2] != 1 or t.shape[3] != D:
+            failed = True
+        elif S % F != 0:
+            raise ValueError(f"Validate failed: S({S}) must be divisible by F({F}).")
+    if failed:
+        raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
+    if t.stride()[-1] != 1:
+        raise ValueError(f"Validate failed: not contiguous on dim D.")
+
+
+def validate_gate(t: Union[torch.Tensor, int], B: int, S: int, D: int):
+    if not isinstance(t, torch.Tensor):
+        return
+    validate_scale_shift(t, B, S, D)
+
+
+@torch.library.custom_op("sglang::fused_norm_scale_shift", mutates_args=())
+def fused_norm_scale_shift(
+    x: torch.Tensor,
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    norm_type: str,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    """
+    Fuse: norm(x) * (1 + scale) + shift
+      where norm is either layernorm or rmsnorm.
+
+    Expects:
+      - x: [B, S, D]
+      - weight/bias: None, [D]
+      - scale/shift: [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D]
+      - norm_type: str, "layer" or "rms"
+      - eps: Optional[float], default: 1e-5
+
+    D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
+    thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
+    """
+    from sglang.jit_kernel.diffusion.norm_scale_shift_native import (
+        try_fused_norm_scale_shift as _try_qwen_native_norm_scale_shift,
+    )
+
+    native_y = _try_qwen_native_norm_scale_shift(
+        x, weight, bias, scale, shift, norm_type, eps
+    )
+    if native_y is not None:
+        return native_y
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    # Tensor Validation
+    BSD = x.shape
+    validate_x(x, *BSD)
+    validate_weight_bias(weight, BSD[-1])
+    validate_weight_bias(bias, BSD[-1])
+    validate_scale_shift(scale, *BSD)
+    validate_scale_shift(shift, *BSD)
+
+    if norm_type == "layer" or norm_type == "rms":
+        D = x.shape[-1]
+        if D % 256 != 0 or D > 8192:
+            raise ValueError(
+                f"D={D} not supported, must be multiple of 256 and <= 8192"
+            )
+        y = torch.empty_like(x)  # create output tensor
+        scale = broadcast_tensor_for_bsfd(scale, *x.shape)  # handle various shapes
+        shift = broadcast_tensor_for_bsfd(shift, *x.shape)  # handle various shapes
+        # Use scalar placeholders for None tensors as a workaround, since the CuTe DSL
+        # TVM-FFI backend does not support None parameters. scalar values do not result
+        # in code generation and have no impact on runtime performance.
+        weight = 1 if weight is None else weight
+        bias = 0 if bias is None else bias
+        ResOut, Residual, Gate = 0, 0, 1
+        torch_tensors = [y, ResOut, Residual, x, Gate, weight, bias, scale, shift]
+        # Compile cache
+        hash_key = ScaleResidualNormScaleShift.make_hash_key(norm_type, *torch_tensors)
+        compiled_fn = _COMPILE_CACHE.get(hash_key)
+        if compiled_fn is None:
+            kernel = ScaleResidualNormScaleShift(D, norm_type)
+            fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
+            compiled_fn = cute.compile(
+                kernel, *fake_sig_args, options="--enable-tvm-ffi"
+            )
+            _COMPILE_CACHE[hash_key] = compiled_fn
+        # Execute
+        compiled_fn(*torch_tensors, eps, stream)
+        return y
+    else:
+        raise ValueError(f'norm_type must be one of "layer" and "rms"')
+
+
+@fused_norm_scale_shift.register_fake
+def _fused_norm_scale_shift_fake(x, weight, bias, scale, shift, norm_type, eps=1e-5):
+    y = x.new_empty(x.shape)
+    return y
+
+
+@torch.library.custom_op(
+    "sglang::fused_scale_residual_norm_scale_shift", mutates_args=()
+)
+def fused_scale_residual_norm_scale_shift(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: Optional[torch.Tensor],  # Union[Optional[torch.Tensor], int] indeed
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    norm_type: str,
+    eps: float = 1e-5,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Fuse: norm(residual + gate * x) * (1 + scale) + shift
+      where norm is either layernorm or rmsnorm.
+
+    Expects:
+      - residual, x: [B, S, D]
+      - gate: None, [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D]
+      - weight/bias: None, [D]
+      - scale/shift: [1], [D], [1/B, D], [1/B, 1/S, D] or [B, F, 1, D]
+      - norm_type: str, "layer" or "rms"
+      - eps: Optional[float], default: 1e-5
+
+    D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
+    thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
+    """
+    from sglang.jit_kernel.diffusion.norm_scale_shift_native import (
+        try_fused_scale_residual_norm_scale_shift as _try_qwen_native_residual_path,
+    )
+
+    native_out = _try_qwen_native_residual_path(
+        residual, x, gate, weight, bias, scale, shift, norm_type, eps
+    )
+    if native_out is not None:
+        return native_out
+    # Tensor Validation
+    BSD = x.shape
+    validate_x(x, *BSD)
+    validate_x(residual, *BSD)
+    validate_gate(gate, *BSD)
+    validate_weight_bias(weight, BSD[-1])
+    validate_weight_bias(bias, BSD[-1])
+    validate_scale_shift(scale, *BSD)
+    validate_scale_shift(shift, *BSD)
+    if norm_type == "layer" or norm_type == "rms":
+        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+        D = x.shape[-1]
+        if D % 256 != 0 or D > 8192:
+            raise ValueError(
+                f"D={D} not supported, must be multiple of 256 and <= 8192"
+            )
+        y = torch.empty_like(x)  # create output tensor
+        resi_out = torch.empty_like(x)  # create output tensor
+        gate = broadcast_tensor_for_bsfd(gate, *x.shape)  # handle various shapes
+        scale = broadcast_tensor_for_bsfd(scale, *x.shape)  # handle various shapes
+        shift = broadcast_tensor_for_bsfd(shift, *x.shape)  # handle various shapes
+        # Use scalar placeholders for None tensors as a workaround, since the CuTe DSL
+        # TVM-FFI backend does not support None parameters. scalar values do not result
+        # in code generation and have no impact on runtime performance.
+        gate = 1 if gate is None else gate
+        weight = 1 if weight is None else weight
+        bias = 0 if bias is None else bias
+        torch_tensors = [y, resi_out, residual, x, gate, weight, bias, scale, shift]
+        # Compile cache
+        hash_key = ScaleResidualNormScaleShift.make_hash_key(norm_type, *torch_tensors)
+        compiled_fn = _COMPILE_CACHE.get(hash_key)
+        if compiled_fn is None:
+            kernel = ScaleResidualNormScaleShift(D, norm_type)
+            fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
+            compiled_fn = cute.compile(
+                kernel, *fake_sig_args, options="--enable-tvm-ffi"
+            )
+            _COMPILE_CACHE[hash_key] = compiled_fn
+        # Execute
+        compiled_fn(*torch_tensors, eps, stream)
+        return y, resi_out
+    else:
+        raise ValueError(f'norm_type must be one of "layer" and "rms"')
+
+
+@fused_scale_residual_norm_scale_shift.register_fake
+def _fused_scale_residual_norm_scale_shift_fake(
+    residual, x, gate, weight, bias, scale, shift, norm_type, eps=1e-5
+):
+    y = x.new_empty(x.shape)
+    residual_out = x.new_empty(x.shape)
+    return y, residual_out

python/sglang/jit_kernel/diffusion/cutedsl/utils.py

@@ -0,0 +1,52 @@
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+import torch
+
+WARP_SIZE = 32
+
+TORCH_TO_CUTE_DTYPE = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+}
+
+
+def to_cute_arg(
+    t,
+    *,
+    assume_aligned: Optional[int] = 32,
+    use_32bit_stride: bool = False,
+    enable_tvm_ffi: bool = True,
+):
+    """
+    Convert a Python value into a CuTeDSL value.
+    """
+    if isinstance(t, torch.Tensor):
+        return cute.runtime.from_dlpack(
+            t,
+            assumed_align=assume_aligned,
+            use_32bit_stride=use_32bit_stride,
+            enable_tvm_ffi=enable_tvm_ffi,
+        )
+    if isinstance(t, int):
+        return cutlass.Int32(t)
+    if isinstance(t, float):
+        return cutlass.Float32(t)
+    return t
+
+
+def to_fake_cute_args(t: torch.Tensor):
+    if isinstance(t, torch.Tensor):
+        # Only keep the last dim as compile-time value to maximum compiled kernel reuse
+        # e.g. (1,2,1536):(3027,1536,1) -> (?,?,1536):(?,?,1)
+        D = t.shape[-1]
+        dtype = TORCH_TO_CUTE_DTYPE[t.dtype]
+        shape = (*(cute.sym_int() for _ in range(t.ndim - 1)), D)
+        stride = (*(cute.sym_int(divisibility=D) for _ in range(t.ndim - 1)), 1)
+        fake_t = cute.runtime.make_fake_tensor(
+            dtype, shape, stride, memspace=cute.AddressSpace.gmem, assumed_align=32
+        )
+        return fake_t
+    return to_cute_arg(t)

python/sglang/jit_kernel/diffusion/flydsl/fused_residual_norm.py

@@ -0,0 +1,902 @@
+"""FlyDSL fused normalization kernels for AMD ROCm (gfx950).
+
+Provides two fused kernels:
+  - flydsl_fused_residual_norm_scale_shift:
+        residual_add + gate_mul + RMSNorm/LayerNorm + scale·shift
+  - flydsl_norm_scale_shift:
+        RMSNorm/LayerNorm + scale·shift
+
+Both kernels use register-cache optimization: Phase 2 (scale·shift)
+reuses f32 intermediate values from Phase 1 (norm) registers instead
+of re-reading from HBM, saving ~20% bandwidth.
+"""
+
+from typing import Optional, Tuple
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+import torch
+from flydsl._mlir import ir
+from flydsl._mlir.dialects import arith as arith_ops
+from flydsl._mlir.dialects import gpu as _gpu
+from flydsl._mlir.dialects import math as math_ops
+from flydsl._mlir.dialects import memref as _memref
+from flydsl._mlir.dialects import (
+    scf,
+)
+from flydsl._mlir.dialects import vector as _vector
+from flydsl.compiler.kernel_function import CompilationContext
+from flydsl.expr import arith, buffer_ops, const_expr, range_constexpr
+from flydsl.expr.arith import ArithValue, CmpIPredicate
+from flydsl.expr.typing import Int32, T
+
+WARP_SIZE = 64
+_VEC = 8
+_NUM_WAVES = 10
+FLYDSL_NORM_MIN_ALIGNED_DIM = WARP_SIZE * _NUM_WAVES * _VEC  # 5120
+
+
+def _v(x):
+    return buffer_ops._unwrap_value(x)
+
+
+def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: bool):
+    VEC = _VEC
+    NUM_WAVES = _NUM_WAVES
+    BLOCK = NUM_WAVES * WARP_SIZE
+    assert (
+        D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0
+    ), f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}"
+    NUM_ITERS = D // (BLOCK * VEC)
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def flydsl_fused_residual_norm_ss_kernel(
+        y_ptr: fx.Tensor,
+        res_out_ptr: fx.Tensor,
+        res_ptr: fx.Tensor,
+        x_ptr: fx.Tensor,
+        gate_ptr: fx.Tensor,
+        weight_ptr: fx.Tensor,
+        bias_ptr: fx.Tensor,
+        scale_ptr: fx.Tensor,
+        shift_ptr: fx.Tensor,
+        total_rows: Int32,
+        gate_stride: Int32,
+        scale_stride: Int32,
+        shift_stride: Int32,
+    ):
+        row = fx.block_idx.x
+        tid = fx.thread_idx.x
+
+        i32 = T.i32
+        f32 = T.f32
+        bf16 = T.bf16
+        vec_f32_t = ir.VectorType.get([VEC], f32)
+        vec_bf16_t = ir.VectorType.get([VEC], bf16)
+
+        y_rsrc = buffer_ops.create_buffer_resource(y_ptr, max_size=True)
+        ro_rsrc = buffer_ops.create_buffer_resource(res_out_ptr, max_size=True)
+        r_rsrc = buffer_ops.create_buffer_resource(res_ptr, max_size=True)
+        x_rsrc = buffer_ops.create_buffer_resource(x_ptr, max_size=True)
+        g_rsrc = buffer_ops.create_buffer_resource(gate_ptr, max_size=True)
+        w_rsrc = buffer_ops.create_buffer_resource(weight_ptr, max_size=True)
+        b_rsrc = buffer_ops.create_buffer_resource(bias_ptr, max_size=True)
+        sc_rsrc = buffer_ops.create_buffer_resource(scale_ptr, max_size=True)
+        sh_rsrc = buffer_ops.create_buffer_resource(shift_ptr, max_size=True)
+
+        row_i32 = ArithValue(row)
+        tid_i32 = ArithValue(tid)
+        D_i32 = arith.constant(D, type=i32)
+        row_off = row_i32 * D_i32
+        gate_row_off = row_i32 * ArithValue(gate_stride)
+        scale_row_off = row_i32 * ArithValue(scale_stride)
+        shift_row_off = row_i32 * ArithValue(shift_stride)
+
+        c_zero_f32 = arith.constant(0.0, type=f32)
+        c_one_f32 = arith.constant(1.0, type=f32)
+        eps_val = arith.constant(1e-6, type=f32)
+        D_float = arith.constant(float(D), type=f32)
+
+        # LDS
+        LDS_SLOTS = NUM_WAVES * 2 + 2
+        ws_attr = ir.Attribute.parse("#gpu.address_space<workgroup>")
+        lds_i8_type = ir.MemRefType.get(
+            [ir.ShapedType.get_dynamic_size()], T.i8, memory_space=ws_attr
+        )
+        lds_f32_type = ir.MemRefType.get([LDS_SLOTS], f32, memory_space=ws_attr)
+        lds_i8 = _gpu.DynamicSharedMemoryOp(lds_i8_type).result
+        byte_zero = arith_ops.ConstantOp(
+            ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), 0)
+        ).result
+        lds = _memref.ViewOp(lds_f32_type, lds_i8, byte_zero, []).result
+
+        lane_id = tid_i32 % arith.constant(WARP_SIZE, type=i32)
+        wave_id = tid_i32 // arith.constant(WARP_SIZE, type=i32)
+
+        # Phase 1: residual + gate*x, accumulate stats, save f32 in registers
+        _saved_ro_f32 = []
+        partial_sum = _v(c_zero_f32)
+        partial_sum_sq = _v(c_zero_f32)
+
+        for it in range_constexpr(NUM_ITERS):
+            col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant(
+                it * BLOCK * VEC, type=i32
+            )
+            off = row_off + col
+
+            r_vec = buffer_ops.buffer_load(r_rsrc, off, vec_width=VEC, dtype=bf16)
+            x_vec = buffer_ops.buffer_load(x_rsrc, off, vec_width=VEC, dtype=bf16)
+            r_f32 = arith_ops.ExtFOp(vec_f32_t, _v(r_vec)).result
+            x_f32 = arith_ops.ExtFOp(vec_f32_t, _v(x_vec)).result
+
+            if const_expr(has_gate):
+                g_off = gate_row_off + col
+                g_vec = buffer_ops.buffer_load(g_rsrc, g_off, vec_width=VEC, dtype=bf16)
+                g_f32 = arith_ops.ExtFOp(vec_f32_t, _v(g_vec)).result
+                gx = arith_ops.MulFOp(g_f32, x_f32).result
+                ro_f32 = arith_ops.AddFOp(r_f32, gx).result
+            else:
+                ro_f32 = arith_ops.AddFOp(r_f32, x_f32).result
+
+            ro_bf16 = arith_ops.TruncFOp(vec_bf16_t, ro_f32).result
+            buffer_ops.buffer_store(ro_bf16, ro_rsrc, off)
+            _saved_ro_f32.append(ro_f32)
+
+            if const_expr(not is_rms):
+                v_sum = _vector.ReductionOp(
+                    f32, _vector.CombiningKind.ADD, ro_f32
+                ).result
+                partial_sum = arith_ops.AddFOp(partial_sum, v_sum).result
+            ro_sq = arith_ops.MulFOp(ro_f32, ro_f32).result
+            v_sum_sq = _vector.ReductionOp(f32, _vector.CombiningKind.ADD, ro_sq).result
+            partial_sum_sq = arith_ops.AddFOp(partial_sum_sq, v_sum_sq).result
+
+        # Intra-wave shuffle
+        width_c = _v(arith.constant(WARP_SIZE, type=i32))
+        w_sum = partial_sum
+        w_sq = partial_sum_sq
+        for sh in [32, 16, 8, 4, 2, 1]:
+            off_sh = _v(arith.constant(sh, type=i32))
+            if const_expr(not is_rms):
+                peer_sum = _gpu.ShuffleOp(
+                    w_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+                ).shuffleResult
+                w_sum = arith_ops.AddFOp(w_sum, peer_sum).result
+            peer_sq = _gpu.ShuffleOp(
+                w_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+            ).shuffleResult
+            w_sq = arith_ops.AddFOp(w_sq, peer_sq).result
+
+        # Cross-wave LDS
+        lane_0 = arith.cmpi(CmpIPredicate.eq, lane_id, arith.constant(0, type=i32))
+        wave_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(wave_id)).result
+
+        _if_lane0 = scf.IfOp(lane_0)
+        with ir.InsertionPoint(_if_lane0.then_block):
+            if const_expr(not is_rms):
+                _memref.StoreOp(w_sum, lds, [wave_idx])
+            sq_slot = arith_ops.AddIOp(
+                wave_idx,
+                arith_ops.ConstantOp(
+                    ir.IndexType.get(),
+                    ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES),
+                ).result,
+            ).result
+            _memref.StoreOp(w_sq, lds, [sq_slot])
+            scf.YieldOp([])
+        _gpu.BarrierOp()
+
+        wave_0 = arith.cmpi(CmpIPredicate.eq, wave_id, arith.constant(0, type=i32))
+        active = arith.andi(
+            wave_0,
+            arith.cmpi(CmpIPredicate.ult, lane_id, arith.constant(NUM_WAVES, type=i32)),
+        )
+        lane_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(lane_id)).result
+        lane_idx_sq = arith_ops.AddIOp(
+            lane_idx,
+            arith_ops.ConstantOp(
+                ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES)
+            ).result,
+        ).result
+
+        if const_expr(is_rms):
+            _if_active = scf.IfOp(active, [f32], has_else=True)
+            with ir.InsertionPoint(_if_active.then_block):
+                sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result
+                scf.YieldOp([sq_val])
+            with ir.InsertionPoint(_if_active.else_block):
+                scf.YieldOp([_v(c_zero_f32)])
+            loaded_sq = _if_active.results[0]
+            loaded_sum = _v(c_zero_f32)
+        else:
+            _if_active = scf.IfOp(active, [f32, f32], has_else=True)
+            with ir.InsertionPoint(_if_active.then_block):
+                s_val = _memref.LoadOp(lds, [lane_idx]).result
+                sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result
+                scf.YieldOp([s_val, sq_val])
+            with ir.InsertionPoint(_if_active.else_block):
+                scf.YieldOp([_v(c_zero_f32), _v(c_zero_f32)])
+            loaded_sum = _if_active.results[0]
+            loaded_sq = _if_active.results[1]
+
+        final_sum = loaded_sum
+        final_sq = loaded_sq
+        for sh in [32, 16, 8, 4, 2, 1]:
+            off_sh = _v(arith.constant(sh, type=i32))
+            if const_expr(not is_rms):
+                ps = _gpu.ShuffleOp(
+                    final_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+                ).shuffleResult
+                final_sum = arith_ops.AddFOp(final_sum, ps).result
+            pq = _gpu.ShuffleOp(
+                final_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+            ).shuffleResult
+            final_sq = arith_ops.AddFOp(final_sq, pq).result
+
+        both_0 = arith.andi(wave_0, lane_0)
+        final_sum_slot = arith_ops.ConstantOp(
+            ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2)
+        ).result
+        final_sq_slot = arith_ops.ConstantOp(
+            ir.IndexType.get(),
+            ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2 + 1),
+        ).result
+        _if_both = scf.IfOp(both_0)
+        with ir.InsertionPoint(_if_both.then_block):
+            if const_expr(not is_rms):
+                _memref.StoreOp(final_sum, lds, [final_sum_slot])
+            _memref.StoreOp(final_sq, lds, [final_sq_slot])
+            scf.YieldOp([])
+        _gpu.BarrierOp()
+
+        if const_expr(not is_rms):
+            total_sum = _memref.LoadOp(lds, [final_sum_slot]).result
+        else:
+            total_sum = _v(c_zero_f32)
+        total_sq = _memref.LoadOp(lds, [final_sq_slot]).result
+
+        # Norm
+        d_f = _v(D_float)
+        eps_v = _v(eps_val)
+        if const_expr(is_rms):
+            var = arith_ops.DivFOp(total_sq, d_f).result
+            var_eps = arith_ops.AddFOp(var, eps_v).result
+            rstd = math_ops.RsqrtOp(var_eps).result
+            mean = _v(c_zero_f32)
+        else:
+            mean = arith_ops.DivFOp(total_sum, d_f).result
+            mean_sq = arith_ops.MulFOp(mean, mean).result
+            var = arith_ops.SubFOp(
+                arith_ops.DivFOp(total_sq, d_f).result, mean_sq
+            ).result
+            var_eps = arith_ops.AddFOp(var, eps_v).result
+            rstd = math_ops.RsqrtOp(var_eps).result
+
+        # Phase 2: normalize using register-cached f32 values (no HBM re-read)
+        mean_splat = _vector.BroadcastOp(vec_f32_t, mean).result
+        rstd_splat = _vector.BroadcastOp(vec_f32_t, rstd).result
+        one_splat = _vector.BroadcastOp(vec_f32_t, _v(c_one_f32)).result
+
+        for it in range_constexpr(NUM_ITERS):
+            col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant(
+                it * BLOCK * VEC, type=i32
+            )
+            off = row_off + col
+
+            ro_f32 = _saved_ro_f32[it]
+
+            if const_expr(is_rms):
+                x_hat = arith_ops.MulFOp(ro_f32, rstd_splat).result
+            else:
+                centered = arith_ops.SubFOp(ro_f32, mean_splat).result
+                x_hat = arith_ops.MulFOp(centered, rstd_splat).result
+
+            if const_expr(has_weight):
+                w_vec = buffer_ops.buffer_load(w_rsrc, col, vec_width=VEC, dtype=bf16)
+                w_f32 = arith_ops.ExtFOp(vec_f32_t, _v(w_vec)).result
+                x_hat = arith_ops.MulFOp(x_hat, w_f32).result
+                b_vec = buffer_ops.buffer_load(b_rsrc, col, vec_width=VEC, dtype=bf16)
+                b_f32 = arith_ops.ExtFOp(vec_f32_t, _v(b_vec)).result
+                x_hat = arith_ops.AddFOp(x_hat, b_f32).result
+
+            sc_off = scale_row_off + col
+            sc_vec = buffer_ops.buffer_load(sc_rsrc, sc_off, vec_width=VEC, dtype=bf16)
+            sc_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sc_vec)).result
+            sc_p1 = arith_ops.AddFOp(one_splat, sc_f32).result
+            x_hat = arith_ops.MulFOp(x_hat, sc_p1).result
+
+            sh_off = shift_row_off + col
+            sh_vec = buffer_ops.buffer_load(sh_rsrc, sh_off, vec_width=VEC, dtype=bf16)
+            sh_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sh_vec)).result
+            y_f32 = arith_ops.AddFOp(x_hat, sh_f32).result
+
+            y_bf16 = arith_ops.TruncFOp(vec_bf16_t, y_f32).result
+            buffer_ops.buffer_store(y_bf16, y_rsrc, off)
+
+    @flyc.jit
+    def launch_fused_norm(
+        y: fx.Tensor,
+        res_out: fx.Tensor,
+        res: fx.Tensor,
+        x: fx.Tensor,
+        gate: fx.Tensor,
+        weight: fx.Tensor,
+        bias: fx.Tensor,
+        scale: fx.Tensor,
+        shift: fx.Tensor,
+        total_rows: fx.Int32,
+        gate_stride: fx.Int32,
+        scale_stride: fx.Int32,
+        shift_stride: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            pass
+        grid_x = arith.index_cast(T.index, total_rows)
+        launcher = flydsl_fused_residual_norm_ss_kernel(
+            y,
+            res_out,
+            res,
+            x,
+            gate,
+            weight,
+            bias,
+            scale,
+            shift,
+            total_rows,
+            gate_stride,
+            scale_stride,
+            shift_stride,
+        )
+        LDS_BYTES = (NUM_WAVES * 2 + 2) * 4
+        launcher.launch(
+            grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), smem=LDS_BYTES, stream=stream
+        )
+
+    return launch_fused_norm
+
+
+_COMPILE_CACHE = {}
+
+
+def _get_or_compile(D, is_rms, has_gate, has_weight, args):
+    key = (D, is_rms, has_gate, has_weight)
+    if key not in _COMPILE_CACHE:
+        launcher = _build_fused_norm_module(D, is_rms, has_gate, has_weight)
+        cf = flyc.compile(launcher, *args)
+        _COMPILE_CACHE[key] = cf
+    return _COMPILE_CACHE[key]
+
+
+def _to_bf16(t):
+    """Convert to bf16 only if not already bf16."""
+    return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16)
+
+
+def _prep_slices(t, B, L, C):
+    """Prepare per-batch tensor slices and kernel row_stride.
+
+    Returns (slices, row_stride) where:
+      slices[b] = tensor to pass to kernel for batch b
+      row_stride = 0 (broadcast: all rows share one row) or C (per-row data)
+    """
+    t = _to_bf16(t)
+
+    if t.numel() < C:
+        row = t.flatten()[0].expand(C).contiguous().unsqueeze(0)
+        return [row] * B, 0
+
+    if t.dim() == 1:
+        return [t.unsqueeze(0).contiguous()] * B, 0
+
+    if t.dim() == 2:
+        if t.shape[0] == 1:
+            return [t.contiguous()] * B, 0
+        return [t.contiguous()] * B, C
+
+    if t.dim() == 3:
+        if t.shape[0] == 1 and t.shape[1] == 1:
+            return [t.reshape(1, C).contiguous()] * B, 0
+        if t.shape[1] == 1:
+            t_c = t.contiguous()
+            return [t_c[b] for b in range(B)], 0
+        t_exp = t.expand(B, L, C).contiguous()
+        return [t_exp[b] for b in range(B)], C
+
+    if t.dim() == 4:
+        nf = t.shape[1]
+        fs = L // nf
+        t_exp = t.expand(B, nf, fs, C).reshape(B, L, C).contiguous()
+        return [t_exp[b] for b in range(B)], C
+
+    t_exp = t.reshape(B, L, C).contiguous()
+    return [t_exp[b] for b in range(B)], C
+
+
+def _ensure_bf16_contig(t):
+    """Return bf16-contiguous view, avoiding copies when possible."""
+    if t.dtype == torch.bfloat16 and t.is_contiguous():
+        return t
+    return _to_bf16(t).contiguous()
+
+
+@torch.library.custom_op(
+    "sglang::flydsl_fused_residual_norm_scale_shift", mutates_args=()
+)
+def flydsl_fused_residual_norm_scale_shift(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: Optional[torch.Tensor],
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    norm_type: str,
+    eps: float = 1e-6,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    B, L, C = x.shape
+    rows = B * L
+    bf16 = torch.bfloat16
+
+    x_2d = _ensure_bf16_contig(x).reshape(rows, C)
+    res_2d = _ensure_bf16_contig(residual).reshape(rows, C)
+    y = torch.empty_like(x_2d)
+    res_out = torch.empty_like(x_2d)
+
+    has_gate = gate is not None
+    if has_gate:
+        g_slices, g_stride = _prep_slices(gate, B, L, C)
+    else:
+        g_slices, g_stride = [x_2d[:1]] * B, 0
+
+    has_weight = weight is not None
+    weight_c = (
+        _ensure_bf16_contig(weight)
+        if has_weight
+        else torch.empty(C, device=x.device, dtype=bf16)
+    )
+    bias_c = (
+        _ensure_bf16_contig(bias)
+        if bias is not None
+        else torch.zeros(C, device=x.device, dtype=bf16)
+    )
+
+    sc_slices, sc_stride = _prep_slices(scale, B, L, C)
+    sh_slices, sh_stride = _prep_slices(shift, B, L, C)
+
+    is_rms = norm_type == "rms"
+    stream = torch.cuda.current_stream()
+
+    dummy_args = (
+        y[:L],
+        res_out[:L],
+        res_2d[:L],
+        x_2d[:L],
+        g_slices[0],
+        weight_c,
+        bias_c,
+        sc_slices[0],
+        sh_slices[0],
+        L,
+        g_stride,
+        sc_stride,
+        sh_stride,
+        stream,
+    )
+    cf = _get_or_compile(C, is_rms, has_gate, has_weight, dummy_args)
+
+    for b in range(B):
+        s, e = b * L, (b + 1) * L
+        cf(
+            y[s:e],
+            res_out[s:e],
+            res_2d[s:e],
+            x_2d[s:e],
+            g_slices[b],
+            weight_c,
+            bias_c,
+            sc_slices[b],
+            sh_slices[b],
+            L,
+            g_stride,
+            sc_stride,
+            sh_stride,
+            stream,
+        )
+
+    return y.view(B, L, C), res_out.view(B, L, C)
+
+
+@flydsl_fused_residual_norm_scale_shift.register_fake
+def _fake_flydsl_fused_residual_norm(
+    residual,
+    x,
+    gate,
+    weight,
+    bias,
+    scale,
+    shift,
+    norm_type,
+    eps=1e-6,
+):
+    B, L, C = x.shape
+    bf16 = torch.bfloat16
+    y = torch.empty(B, L, C, device=x.device, dtype=bf16)
+    res_out = torch.empty(B, L, C, device=x.device, dtype=bf16)
+    return y, res_out
+
+
+###############################################################################
+# _NormScaleShift kernel: norm(x) * (1+scale) + shift  (no residual path)
+###############################################################################
+
+
+def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool):
+    VEC = _VEC
+    NUM_WAVES = _NUM_WAVES
+    BLOCK = NUM_WAVES * WARP_SIZE
+    assert (
+        D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0
+    ), f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}"
+    NUM_ITERS = D // (BLOCK * VEC)
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def flydsl_norm_scale_shift_kernel(
+        y_ptr: fx.Tensor,
+        x_ptr: fx.Tensor,
+        weight_ptr: fx.Tensor,
+        bias_ptr: fx.Tensor,
+        scale_ptr: fx.Tensor,
+        shift_ptr: fx.Tensor,
+        total_rows: Int32,
+        scale_stride: Int32,
+        shift_stride: Int32,
+    ):
+        row = fx.block_idx.x
+        tid = fx.thread_idx.x
+
+        i32 = T.i32
+        f32 = T.f32
+        bf16 = T.bf16
+        vec_f32_t = ir.VectorType.get([VEC], f32)
+        vec_bf16_t = ir.VectorType.get([VEC], bf16)
+
+        y_rsrc = buffer_ops.create_buffer_resource(y_ptr, max_size=True)
+        x_rsrc = buffer_ops.create_buffer_resource(x_ptr, max_size=True)
+        w_rsrc = buffer_ops.create_buffer_resource(weight_ptr, max_size=True)
+        b_rsrc = buffer_ops.create_buffer_resource(bias_ptr, max_size=True)
+        sc_rsrc = buffer_ops.create_buffer_resource(scale_ptr, max_size=True)
+        sh_rsrc = buffer_ops.create_buffer_resource(shift_ptr, max_size=True)
+
+        row_i32 = ArithValue(row)
+        tid_i32 = ArithValue(tid)
+        D_i32 = arith.constant(D, type=i32)
+        row_off = row_i32 * D_i32
+        scale_row_off = row_i32 * ArithValue(scale_stride)
+        shift_row_off = row_i32 * ArithValue(shift_stride)
+
+        c_zero_f32 = arith.constant(0.0, type=f32)
+        c_one_f32 = arith.constant(1.0, type=f32)
+        eps_val = arith.constant(1e-6, type=f32)
+        D_float = arith.constant(float(D), type=f32)
+
+        LDS_SLOTS = NUM_WAVES * 2 + 2
+        ws_attr = ir.Attribute.parse("#gpu.address_space<workgroup>")
+        lds_i8_type = ir.MemRefType.get(
+            [ir.ShapedType.get_dynamic_size()], T.i8, memory_space=ws_attr
+        )
+        lds_f32_type = ir.MemRefType.get([LDS_SLOTS], f32, memory_space=ws_attr)
+        lds_i8 = _gpu.DynamicSharedMemoryOp(lds_i8_type).result
+        byte_zero = arith_ops.ConstantOp(
+            ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), 0)
+        ).result
+        lds = _memref.ViewOp(lds_f32_type, lds_i8, byte_zero, []).result
+
+        lane_id = tid_i32 % arith.constant(WARP_SIZE, type=i32)
+        wave_id = tid_i32 // arith.constant(WARP_SIZE, type=i32)
+
+        # Phase 1: load x, accumulate stats, save f32 in registers
+        _saved_x_f32 = []
+        partial_sum = _v(c_zero_f32)
+        partial_sum_sq = _v(c_zero_f32)
+
+        for it in range_constexpr(NUM_ITERS):
+            col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant(
+                it * BLOCK * VEC, type=i32
+            )
+            off = row_off + col
+
+            x_vec = buffer_ops.buffer_load(x_rsrc, off, vec_width=VEC, dtype=bf16)
+            x_f32 = arith_ops.ExtFOp(vec_f32_t, _v(x_vec)).result
+            _saved_x_f32.append(x_f32)
+
+            if const_expr(not is_rms):
+                v_sum = _vector.ReductionOp(
+                    f32, _vector.CombiningKind.ADD, x_f32
+                ).result
+                partial_sum = arith_ops.AddFOp(partial_sum, v_sum).result
+            x_sq = arith_ops.MulFOp(x_f32, x_f32).result
+            v_sum_sq = _vector.ReductionOp(f32, _vector.CombiningKind.ADD, x_sq).result
+            partial_sum_sq = arith_ops.AddFOp(partial_sum_sq, v_sum_sq).result
+
+        # Intra-wave shuffle reduction
+        width_c = _v(arith.constant(WARP_SIZE, type=i32))
+        w_sum = partial_sum
+        w_sq = partial_sum_sq
+        for sh in [32, 16, 8, 4, 2, 1]:
+            off_sh = _v(arith.constant(sh, type=i32))
+            if const_expr(not is_rms):
+                peer_sum = _gpu.ShuffleOp(
+                    w_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+                ).shuffleResult
+                w_sum = arith_ops.AddFOp(w_sum, peer_sum).result
+            peer_sq = _gpu.ShuffleOp(
+                w_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+            ).shuffleResult
+            w_sq = arith_ops.AddFOp(w_sq, peer_sq).result
+
+        # Cross-wave LDS reduction
+        lane_0 = arith.cmpi(CmpIPredicate.eq, lane_id, arith.constant(0, type=i32))
+        wave_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(wave_id)).result
+
+        _if_lane0 = scf.IfOp(lane_0)
+        with ir.InsertionPoint(_if_lane0.then_block):
+            if const_expr(not is_rms):
+                _memref.StoreOp(w_sum, lds, [wave_idx])
+            sq_slot = arith_ops.AddIOp(
+                wave_idx,
+                arith_ops.ConstantOp(
+                    ir.IndexType.get(),
+                    ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES),
+                ).result,
+            ).result
+            _memref.StoreOp(w_sq, lds, [sq_slot])
+            scf.YieldOp([])
+        _gpu.BarrierOp()
+
+        wave_0 = arith.cmpi(CmpIPredicate.eq, wave_id, arith.constant(0, type=i32))
+        active = arith.andi(
+            wave_0,
+            arith.cmpi(CmpIPredicate.ult, lane_id, arith.constant(NUM_WAVES, type=i32)),
+        )
+        lane_idx = arith_ops.IndexCastOp(ir.IndexType.get(), _v(lane_id)).result
+        lane_idx_sq = arith_ops.AddIOp(
+            lane_idx,
+            arith_ops.ConstantOp(
+                ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES)
+            ).result,
+        ).result
+
+        if const_expr(is_rms):
+            _if_active = scf.IfOp(active, [f32], has_else=True)
+            with ir.InsertionPoint(_if_active.then_block):
+                sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result
+                scf.YieldOp([sq_val])
+            with ir.InsertionPoint(_if_active.else_block):
+                scf.YieldOp([_v(c_zero_f32)])
+            loaded_sq = _if_active.results[0]
+            loaded_sum = _v(c_zero_f32)
+        else:
+            _if_active = scf.IfOp(active, [f32, f32], has_else=True)
+            with ir.InsertionPoint(_if_active.then_block):
+                s_val = _memref.LoadOp(lds, [lane_idx]).result
+                sq_val = _memref.LoadOp(lds, [lane_idx_sq]).result
+                scf.YieldOp([s_val, sq_val])
+            with ir.InsertionPoint(_if_active.else_block):
+                scf.YieldOp([_v(c_zero_f32), _v(c_zero_f32)])
+            loaded_sum = _if_active.results[0]
+            loaded_sq = _if_active.results[1]
+
+        final_sum = loaded_sum
+        final_sq = loaded_sq
+        for sh in [32, 16, 8, 4, 2, 1]:
+            off_sh = _v(arith.constant(sh, type=i32))
+            if const_expr(not is_rms):
+                ps = _gpu.ShuffleOp(
+                    final_sum, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+                ).shuffleResult
+                final_sum = arith_ops.AddFOp(final_sum, ps).result
+            pq = _gpu.ShuffleOp(
+                final_sq, off_sh, width_c, mode=_gpu.ShuffleMode.XOR
+            ).shuffleResult
+            final_sq = arith_ops.AddFOp(final_sq, pq).result
+
+        both_0 = arith.andi(wave_0, lane_0)
+        final_sum_slot = arith_ops.ConstantOp(
+            ir.IndexType.get(), ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2)
+        ).result
+        final_sq_slot = arith_ops.ConstantOp(
+            ir.IndexType.get(),
+            ir.IntegerAttr.get(ir.IndexType.get(), NUM_WAVES * 2 + 1),
+        ).result
+        _if_both = scf.IfOp(both_0)
+        with ir.InsertionPoint(_if_both.then_block):
+            if const_expr(not is_rms):
+                _memref.StoreOp(final_sum, lds, [final_sum_slot])
+            _memref.StoreOp(final_sq, lds, [final_sq_slot])
+            scf.YieldOp([])
+        _gpu.BarrierOp()
+
+        if const_expr(not is_rms):
+            total_sum = _memref.LoadOp(lds, [final_sum_slot]).result
+        else:
+            total_sum = _v(c_zero_f32)
+        total_sq = _memref.LoadOp(lds, [final_sq_slot]).result
+
+        d_f = _v(D_float)
+        eps_v = _v(eps_val)
+        if const_expr(is_rms):
+            var = arith_ops.DivFOp(total_sq, d_f).result
+            var_eps = arith_ops.AddFOp(var, eps_v).result
+            rstd = math_ops.RsqrtOp(var_eps).result
+            mean = _v(c_zero_f32)
+        else:
+            mean = arith_ops.DivFOp(total_sum, d_f).result
+            mean_sq = arith_ops.MulFOp(mean, mean).result
+            var = arith_ops.SubFOp(
+                arith_ops.DivFOp(total_sq, d_f).result, mean_sq
+            ).result
+            var_eps = arith_ops.AddFOp(var, eps_v).result
+            rstd = math_ops.RsqrtOp(var_eps).result
+
+        # Phase 2: normalize from register cache + scale_shift → single output
+        mean_splat = _vector.BroadcastOp(vec_f32_t, mean).result
+        rstd_splat = _vector.BroadcastOp(vec_f32_t, rstd).result
+        one_splat = _vector.BroadcastOp(vec_f32_t, _v(c_one_f32)).result
+
+        for it in range_constexpr(NUM_ITERS):
+            col = tid_i32 * arith.constant(VEC, type=i32) + arith.constant(
+                it * BLOCK * VEC, type=i32
+            )
+            off = row_off + col
+
+            x_f32 = _saved_x_f32[it]
+
+            if const_expr(is_rms):
+                x_hat = arith_ops.MulFOp(x_f32, rstd_splat).result
+            else:
+                centered = arith_ops.SubFOp(x_f32, mean_splat).result
+                x_hat = arith_ops.MulFOp(centered, rstd_splat).result
+
+            if const_expr(has_weight):
+                w_vec = buffer_ops.buffer_load(w_rsrc, col, vec_width=VEC, dtype=bf16)
+                w_f32 = arith_ops.ExtFOp(vec_f32_t, _v(w_vec)).result
+                x_hat = arith_ops.MulFOp(x_hat, w_f32).result
+                b_vec = buffer_ops.buffer_load(b_rsrc, col, vec_width=VEC, dtype=bf16)
+                b_f32 = arith_ops.ExtFOp(vec_f32_t, _v(b_vec)).result
+                x_hat = arith_ops.AddFOp(x_hat, b_f32).result
+
+            sc_off = scale_row_off + col
+            sc_vec = buffer_ops.buffer_load(sc_rsrc, sc_off, vec_width=VEC, dtype=bf16)
+            sc_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sc_vec)).result
+            sc_p1 = arith_ops.AddFOp(one_splat, sc_f32).result
+            x_hat = arith_ops.MulFOp(x_hat, sc_p1).result
+
+            sh_off = shift_row_off + col
+            sh_vec = buffer_ops.buffer_load(sh_rsrc, sh_off, vec_width=VEC, dtype=bf16)
+            sh_f32 = arith_ops.ExtFOp(vec_f32_t, _v(sh_vec)).result
+            y_f32 = arith_ops.AddFOp(x_hat, sh_f32).result
+
+            y_bf16 = arith_ops.TruncFOp(vec_bf16_t, y_f32).result
+            buffer_ops.buffer_store(y_bf16, y_rsrc, off)
+
+    @flyc.jit
+    def launch_norm_ss(
+        y: fx.Tensor,
+        x: fx.Tensor,
+        weight: fx.Tensor,
+        bias: fx.Tensor,
+        scale: fx.Tensor,
+        shift: fx.Tensor,
+        total_rows: fx.Int32,
+        scale_stride: fx.Int32,
+        shift_stride: fx.Int32,
+        stream: fx.Stream = fx.Stream(None),
+    ):
+        ctx = CompilationContext.get_current()
+        with ir.InsertionPoint(ctx.gpu_module_body):
+            pass
+        grid_x = arith.index_cast(T.index, total_rows)
+        launcher = flydsl_norm_scale_shift_kernel(
+            y,
+            x,
+            weight,
+            bias,
+            scale,
+            shift,
+            total_rows,
+            scale_stride,
+            shift_stride,
+        )
+        LDS_BYTES = (NUM_WAVES * 2 + 2) * 4
+        launcher.launch(
+            grid=(grid_x, 1, 1), block=(BLOCK, 1, 1), smem=LDS_BYTES, stream=stream
+        )
+
+    return launch_norm_ss
+
+
+_NSS_COMPILE_CACHE = {}
+
+
+def _get_or_compile_nss(D, is_rms, has_weight, args):
+    key = ("nss", D, is_rms, has_weight)
+    if key not in _NSS_COMPILE_CACHE:
+        launcher = _build_norm_scale_shift_module(D, is_rms, has_weight)
+        cf = flyc.compile(launcher, *args)
+        _NSS_COMPILE_CACHE[key] = cf
+    return _NSS_COMPILE_CACHE[key]
+
+
+@torch.library.custom_op("sglang::flydsl_norm_scale_shift", mutates_args=())
+def flydsl_norm_scale_shift(
+    x: torch.Tensor,
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    norm_type: str,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    B, L, C = x.shape
+    rows = B * L
+    bf16 = torch.bfloat16
+
+    x_2d = _ensure_bf16_contig(x).reshape(rows, C)
+    y = torch.empty_like(x_2d)
+
+    has_weight = weight is not None
+    weight_c = (
+        _ensure_bf16_contig(weight)
+        if has_weight
+        else torch.empty(C, device=x.device, dtype=bf16)
+    )
+    bias_c = (
+        _ensure_bf16_contig(bias)
+        if bias is not None
+        else torch.zeros(C, device=x.device, dtype=bf16)
+    )
+
+    sc_slices, sc_stride = _prep_slices(scale, B, L, C)
+    sh_slices, sh_stride = _prep_slices(shift, B, L, C)
+
+    is_rms = norm_type == "rms"
+    stream = torch.cuda.current_stream()
+
+    dummy_args = (
+        y[:L],
+        x_2d[:L],
+        weight_c,
+        bias_c,
+        sc_slices[0],
+        sh_slices[0],
+        L,
+        sc_stride,
+        sh_stride,
+        stream,
+    )
+    cf = _get_or_compile_nss(C, is_rms, has_weight, dummy_args)
+
+    for b in range(B):
+        s, e = b * L, (b + 1) * L
+        cf(
+            y[s:e],
+            x_2d[s:e],
+            weight_c,
+            bias_c,
+            sc_slices[b],
+            sh_slices[b],
+            L,
+            sc_stride,
+            sh_stride,
+            stream,
+        )
+
+    return y.view(B, L, C)
+
+
+@flydsl_norm_scale_shift.register_fake
+def _fake_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps=1e-6):
+    B, L, C = x.shape
+    return torch.empty(B, L, C, device=x.device, dtype=torch.bfloat16)

python/sglang/jit_kernel/diffusion/group_norm_silu.py

@@ -0,0 +1,35 @@
+import torch
+from torch import nn
+
+
+def apply_group_norm_silu(
+    x: torch.Tensor,
+    norm: nn.Module,
+    activation: nn.Module,
+) -> torch.Tensor:
+    if (
+        x.is_cuda
+        and not torch.is_grad_enabled()
+        and not x.requires_grad
+        and isinstance(norm, nn.GroupNorm)
+        and isinstance(activation, nn.SiLU)
+        and not activation.inplace
+        and norm.affine
+        and norm.weight is not None
+        and norm.bias is not None
+    ):
+        from sglang.jit_kernel.diffusion.triton.group_norm_silu import (
+            triton_group_norm_silu,
+        )
+
+        return triton_group_norm_silu(
+            x,
+            norm.weight,
+            norm.bias,
+            num_groups=norm.num_groups,
+            eps=norm.eps,
+        )
+    return activation(norm(x))
+
+
+__all__ = ["apply_group_norm_silu"]

python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_ltx2_qknorm_split_rope_module() -> Module:
+    return load_jit(
+        "diffusion_ltx2_qknorm_split_rope",
+        cuda_files=["diffusion/ltx2_qknorm_split_rope.cuh"],
+        cuda_wrappers=[
+            (
+                "ltx2_qknorm_split_rope_pair",
+                "sglang_ltx2_qknorm_split_rope::LTX2QKNormSplitRopeKernel::run",
+            )
+        ],
+    )
+
+
+def _fake_impl(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+    num_heads: int,
+    head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return torch.empty_like(q, dtype=torch.bfloat16), torch.empty_like(
+        k, dtype=torch.bfloat16
+    )
+
+
+@register_custom_op(
+    op_name="diffusion_ltx2_qknorm_split_rope",
+    mutates_args=[],
+    fake_impl=_fake_impl,
+)
+def _ltx2_qknorm_split_rope_custom_op(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float,
+    num_heads: int,
+    head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    q_out = torch.empty_like(q, dtype=torch.bfloat16)
+    k_out = torch.empty_like(k, dtype=torch.bfloat16)
+    module = _jit_ltx2_qknorm_split_rope_module()
+    module.ltx2_qknorm_split_rope_pair(
+        q_out,
+        k_out,
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        float(eps),
+        int(num_heads),
+        int(head_dim),
+    )
+    return q_out, k_out
+
+
+def _supported_side(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    num_heads: int,
+    head_dim: int,
+) -> bool:
+    return (
+        x.is_cuda
+        and cos.is_cuda
+        and sin.is_cuda
+        and weight.is_cuda
+        and x.device == cos.device == sin.device == weight.device
+        and x.dtype == torch.bfloat16
+        and cos.dtype == torch.bfloat16
+        and sin.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and x.ndim == 3
+        and cos.ndim == 4
+        and sin.ndim == 4
+        and x.is_contiguous()
+        and cos.shape == sin.shape
+        and cos.shape[0] == x.shape[0]
+        and cos.shape[1] == num_heads
+        and cos.shape[2] == x.shape[1]
+        and cos.shape[3] * 2 == head_dim
+        and x.shape[2] == num_heads * head_dim
+        and x.shape[2] == weight.shape[0]
+        and weight.ndim == 1
+        and head_dim % 2 == 0
+        and x.shape[2] % 4 == 0
+        and cos.stride(-1) == 1
+        and sin.stride(-1) == 1
+    )
+
+
+def _is_sm100_or_newer(x: torch.Tensor) -> bool:
+    if not x.is_cuda:
+        return False
+    try:
+        return torch.cuda.get_device_capability(x.device)[0] >= 10
+    except RuntimeError:
+        return False
+
+
+def can_use_ltx2_qknorm_split_rope_cuda(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    *,
+    num_heads: int,
+    head_dim: int,
+) -> bool:
+    return (
+        _is_sm100_or_newer(q)
+        and _supported_side(
+            q,
+            q_cos,
+            q_sin,
+            q_weight,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+        and _supported_side(
+            k,
+            k_cos,
+            k_sin,
+            k_weight,
+            num_heads=num_heads,
+            head_dim=head_dim,
+        )
+    )
+
+
+def ltx2_qknorm_split_rope_cuda(
+    q: torch.Tensor,
+    q_cos: torch.Tensor,
+    q_sin: torch.Tensor,
+    q_weight: torch.Tensor,
+    k: torch.Tensor,
+    k_cos: torch.Tensor,
+    k_sin: torch.Tensor,
+    k_weight: torch.Tensor,
+    *,
+    eps: float,
+    num_heads: int,
+    head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not can_use_ltx2_qknorm_split_rope_cuda(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        num_heads=num_heads,
+        head_dim=head_dim,
+    ):
+        raise RuntimeError("unsupported input for LTX2 QKNorm split-RoPE CUDA")
+    return _ltx2_qknorm_split_rope_custom_op(
+        q,
+        q_cos,
+        q_sin,
+        q_weight,
+        k,
+        k_cos,
+        k_sin,
+        k_weight,
+        float(eps),
+        int(num_heads),
+        int(head_dim),
+    )

python/sglang/jit_kernel/diffusion/norm_scale_shift_native.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+_HIDDEN = 3072
+_ALIGN = 32
+
+
+def _aligned(t: torch.Tensor) -> bool:
+    return t.data_ptr() % _ALIGN == 0
+
+
+def _blackwell_or_newer(device: torch.device) -> bool:
+    return (
+        torch.cuda.is_available() and torch.cuda.get_device_capability(device)[0] >= 10
+    )
+
+
+def _qwen_activation(t, like=None) -> bool:
+    return (
+        isinstance(t, torch.Tensor)
+        and t.is_cuda
+        and t.dtype == torch.bfloat16
+        and t.ndim == 3
+        and t.shape[0] == 1
+        and t.shape[-1] == _HIDDEN
+        and t.numel() > 0
+        and t.is_contiguous()
+        and _aligned(t)
+        and (like is None or (t.device == like.device and t.shape == like.shape))
+    )
+
+
+def _row_bf16(t, device: torch.device):
+    if (
+        not isinstance(t, torch.Tensor)
+        or t.dtype != torch.bfloat16
+        or not t.is_cuda
+        or t.device != device
+        or t.ndim < 1
+        or t.stride(-1) != 1
+    ):
+        return None
+    if t.shape == (_HIDDEN,):
+        row = t
+    elif t.shape in ((1, _HIDDEN), (1, 1, _HIDDEN)):
+        row = t.reshape(_HIDDEN)
+    else:
+        return None
+    return row if _aligned(row) else None
+
+
+@cache_once
+def _jit_norm_scale_shift_module() -> Module:
+    return load_jit(
+        "qwen_image_norm_scale_shift_native",
+        cuda_files=["diffusion/norm_scale_shift.cuh"],
+        cuda_wrappers=[
+            (
+                "qwen_image_nss_bf16_row",
+                "sglang_norm_scale_shift::QwenImageNormScaleShiftKernel::run",
+            ),
+            (
+                "qwen_image_srnss_bf16_row",
+                "sglang_norm_scale_shift::"
+                "QwenImageScaleResidualNormScaleShiftKernel::run",
+            ),
+        ],
+    )
+
+
+_module = _jit_norm_scale_shift_module
+
+
+def try_fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps):
+    if norm_type != "layer" or weight is not None or bias is not None:
+        return None
+    if not _qwen_activation(x) or not _blackwell_or_newer(x.device):
+        return None
+
+    scale = _row_bf16(scale, x.device)
+    shift = _row_bf16(shift, x.device)
+    if scale is None or shift is None:
+        return None
+
+    y = torch.empty_like(x)
+    _module().qwen_image_nss_bf16_row(
+        y.view(-1, _HIDDEN), x.view(-1, _HIDDEN), scale, shift, float(eps)
+    )
+    return y
+
+
+def try_fused_scale_residual_norm_scale_shift(
+    residual, x, gate, weight, bias, scale, shift, norm_type, eps
+):
+    if norm_type != "layer" or weight is not None or bias is not None:
+        return None
+    if not (
+        _qwen_activation(x)
+        and _qwen_activation(residual, x)
+        and _blackwell_or_newer(x.device)
+    ):
+        return None
+
+    gate = _row_bf16(gate, x.device)
+    scale = _row_bf16(scale, x.device)
+    shift = _row_bf16(shift, x.device)
+    if gate is None or scale is None or shift is None:
+        return None
+
+    y = torch.empty_like(x)
+    residual_out = torch.empty_like(x)
+    _module().qwen_image_srnss_bf16_row(
+        y.view(-1, _HIDDEN),
+        residual_out.view(-1, _HIDDEN),
+        residual.view(-1, _HIDDEN),
+        x.view(-1, _HIDDEN),
+        gate,
+        scale,
+        shift,
+        float(eps),
+    )
+    return y, residual_out

python/sglang/jit_kernel/diffusion/qknorm_rope.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+logger = logging.getLogger(__name__)
+
+
+@cache_once
+def _jit_qknorm_rope_module(
+    head_dim: int,
+    rope_dim: int,
+    is_neox: bool,
+    dtype: torch.dtype,
+) -> Module:
+    args = make_cpp_args(head_dim, rope_dim, is_neox, is_arch_support_pdl(), dtype)
+    return load_jit(
+        "qknorm_rope",
+        *args,
+        cuda_files=["diffusion/qknorm_rope.cuh"],
+        cuda_wrappers=[("qknorm_rope", f"QKNormRopeKernel<{args}>::run")],
+    )
+
+
+@torch.compiler.assume_constant_result
+@cache_once
+def can_use_fused_inplace_qknorm_rope(
+    head_dim: int,
+    rope_dim: int,
+    is_neox: bool,
+    dtype: torch.dtype,
+) -> bool:
+    if head_dim not in (64, 128, 256):
+        logger.warning(f"Unsupported head_dim={head_dim} for JIT fused QKNorm+RoPE")
+        return False
+    if rope_dim <= 0 or rope_dim > head_dim:
+        logger.warning(
+            f"Unsupported rope_dim={rope_dim} for head_dim={head_dim} in fused QKNorm+RoPE"
+        )
+        return False
+    elems_per_thread = head_dim // 32
+    if rope_dim % elems_per_thread != 0:
+        logger.warning(
+            "rope_dim=%s must be divisible by per-thread width=%s for fused QKNorm+RoPE",
+            rope_dim,
+            elems_per_thread,
+        )
+        return False
+    if is_neox:
+        rotary_lanes = rope_dim // elems_per_thread
+        if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
+            logger.warning(
+                "rope_dim=%s yields invalid rotary_lanes=%s for neox fused QKNorm+RoPE; rotary lane count must be a power of 2",
+                rope_dim,
+                rotary_lanes,
+            )
+            return False
+    try:
+        _jit_qknorm_rope_module(head_dim, rope_dim, is_neox, dtype)
+        return True
+    except Exception as e:
+        logger.warning(f"Failed to load JIT fused QKNorm+RoPE kernel: {e}")
+        return False
+
+
+@register_custom_op(mutates_args=["q", "k"])
+def fused_inplace_qknorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    *,
+    is_neox: bool,
+    eps: float = 1e-6,
+    head_dim: int = 0,
+    rope_dim: int = 0,
+) -> None:
+    head_dim = head_dim or q.size(-1)
+    rope_dim = rope_dim or cos_sin_cache.size(-1)
+    module = _jit_qknorm_rope_module(head_dim, rope_dim, is_neox, q.dtype)
+    module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps)

python/sglang/jit_kernel/diffusion/residual_gate_add.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)
+
+
+@cache_once
+def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "diffusion_residual_gate_add",
+        *args,
+        cuda_files=["diffusion/residual_gate_add.cuh"],
+        cuda_wrappers=[
+            (
+                "residual_gate_add",
+                "sglang_residual_gate_add::" f"ResidualGateAddKernel<{args}>::run",
+            ),
+        ],
+    )
+
+
+def _fake_impl(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> torch.Tensor:
+    return torch.empty_like(residual)
+
+
+@register_custom_op(
+    op_name="diffusion_residual_gate_add",
+    mutates_args=[],
+    fake_impl=_fake_impl,
+)
+def _residual_gate_add_custom_op(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> torch.Tensor:
+    out = torch.empty_like(residual)
+    module = _jit_residual_gate_add_module(residual.dtype)
+    module.residual_gate_add(out, residual, update, gate)
+    return out
+
+
+def _is_row_broadcast_gate(residual: torch.Tensor, gate: torch.Tensor) -> bool:
+    if gate.dim() != residual.dim() or gate.shape[-1] != residual.shape[-1]:
+        return False
+    row_dim = gate.dim() - 2
+    return gate.shape[row_dim] == 1 and all(size == 1 for size in gate.shape[:-1])
+
+
+def can_use_residual_gate_add_cuda(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> bool:
+    return (
+        residual.dtype in _SUPPORTED_DTYPES
+        and residual.dtype == update.dtype
+        and residual.dtype == gate.dtype
+        and residual.is_cuda
+        and update.is_cuda
+        and gate.is_cuda
+        and residual.device == update.device == gate.device
+        and residual.dim() >= 2
+        and update.shape == residual.shape
+        and (gate.shape == residual.shape or _is_row_broadcast_gate(residual, gate))
+        and residual.is_contiguous()
+        and update.is_contiguous()
+        and gate.is_contiguous()
+    )
+
+
+def residual_gate_add_cuda(
+    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
+) -> torch.Tensor:
+    if not can_use_residual_gate_add_cuda(residual, update, gate):
+        raise RuntimeError("unsupported input for residual_gate_add CUDA")
+    return _residual_gate_add_custom_op(residual, update, gate)

python/sglang/jit_kernel/diffusion/triton/causal_conv3d_pad.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+
+@triton.jit
+def _fused_cat_pad_5d_kernel(
+    x_ptr,
+    cache_ptr,
+    out_ptr,
+    total,
+    channels,
+    t_size,
+    h_size,
+    w_size,
+    cache_t,
+    out_t,
+    out_h,
+    out_w,
+    pad_d_left,
+    pad_h_top,
+    pad_w_left,
+    block_size: tl.constexpr,
+):
+    offsets = tl.program_id(0) * block_size + tl.arange(0, block_size)
+    mask = offsets < total
+
+    ow = offsets % out_w
+    tmp = offsets // out_w
+    oh = tmp % out_h
+    tmp = tmp // out_h
+    out = tmp % out_t
+    tmp = tmp // out_t
+    oc = tmp % channels
+    ob = tmp // channels
+
+    iw = ow - pad_w_left
+    ih = oh - pad_h_top
+    src_t = out - pad_d_left
+
+    valid = (
+        mask
+        & (iw >= 0)
+        & (iw < w_size)
+        & (ih >= 0)
+        & (ih < h_size)
+        & (src_t >= 0)
+        & (src_t < cache_t + t_size)
+    )
+    from_cache = src_t < cache_t
+
+    x_t = src_t - cache_t
+    clamped_iw = tl.minimum(tl.maximum(iw, 0), w_size - 1)
+    clamped_ih = tl.minimum(tl.maximum(ih, 0), h_size - 1)
+    clamped_x_t = tl.minimum(tl.maximum(x_t, 0), t_size - 1)
+    clamped_src_t = tl.minimum(tl.maximum(src_t, 0), cache_t - 1)
+
+    x_offsets = (
+        ((ob * channels + oc) * t_size + clamped_x_t) * h_size + clamped_ih
+    ) * w_size + clamped_iw
+    cache_offsets = (
+        ((ob * channels + oc) * cache_t + clamped_src_t) * h_size + clamped_ih
+    ) * w_size + clamped_iw
+
+    x_vals = tl.load(x_ptr + x_offsets, mask=valid & ~from_cache, other=0.0)
+    cache_vals = tl.load(cache_ptr + cache_offsets, mask=valid & from_cache, other=0.0)
+    vals = tl.where(from_cache, cache_vals, x_vals)
+    tl.store(out_ptr + offsets, vals, mask=mask)
+
+
+def fused_causal_conv3d_cat_pad(
+    x: torch.Tensor,
+    cache_x: torch.Tensor,
+    padding: list[int] | tuple[int, ...],
+) -> torch.Tensor:
+    width_left, width_right, height_top, height_bottom, depth_left, depth_right = (
+        padding
+    )
+    depth_left -= cache_x.shape[2]
+    assert depth_left >= 0
+    assert depth_right == 0
+    assert width_left == width_right
+    assert height_top == height_bottom
+
+    bsz, channels, t_size, h_size, w_size = x.shape
+    cache_t = cache_x.shape[2]
+    out = torch.empty(
+        (
+            bsz,
+            channels,
+            t_size + cache_t + depth_left + depth_right,
+            h_size + height_top + height_bottom,
+            w_size + width_left + width_right,
+        ),
+        device=x.device,
+        dtype=x.dtype,
+    )
+    block_size = 256
+    total = out.numel()
+    grid = (triton.cdiv(total, block_size),)
+    with torch.get_device_module().device(x.device):
+        _fused_cat_pad_5d_kernel[grid](
+            x,
+            cache_x,
+            out,
+            total,
+            channels,
+            t_size,
+            h_size,
+            w_size,
+            cache_t,
+            out.shape[2],
+            out.shape[3],
+            out.shape[4],
+            depth_left,
+            height_top,
+            width_left,
+            block_size,
+        )
+    return out

python/sglang/jit_kernel/diffusion/triton/group_norm_silu.py

@@ -0,0 +1,412 @@
+import math
+
+import torch
+import torch.nn.functional as F
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.srt.utils.custom_op import register_custom_op
+
+_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16, torch.float32}
+_LARGE_GROUP_THRESHOLD = 1 << 18
+_BLOCK_SIZE = 4096
+_BLOCKS_PER_PROGRAM = 2
+_CHUNK_SIZE = _BLOCK_SIZE * _BLOCKS_PER_PROGRAM
+
+
+@triton.jit
+def _group_norm_silu_contiguous_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    channels,
+    spatial_size,
+    channels_per_group,
+    group_size,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    group_id = tl.program_id(0).to(tl.int64)
+    batch_id = tl.program_id(1).to(tl.int64)
+
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    for off in range(0, group_size, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        sum_val += tl.sum(x, axis=0)
+        sum_sq += tl.sum(x * x, axis=0)
+
+    inv_group = 1.0 / group_size
+    mean = sum_val * inv_group
+    var = sum_sq * inv_group - mean * mean
+    rstd = tl.rsqrt(var + eps)
+
+    weight_group_offset = group_id * channels_per_group
+    for off in range(0, group_size, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        channel_offsets = weight_group_offset + idx // spatial_size
+        weight = tl.load(weight_ptr + channel_offsets, mask=mask, other=1.0).to(
+            tl.float32
+        )
+        bias = tl.load(bias_ptr + channel_offsets, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+@triton.jit
+def _group_norm_stats_kernel(
+    input_ptr,
+    partial_sum_ptr,
+    partial_sq_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        sum_val += tl.sum(x, axis=0)
+        sum_sq += tl.sum(x * x, axis=0)
+
+    partial_index = row * chunks_per_row + chunk_id
+    tl.store(partial_sum_ptr + partial_index, sum_val)
+    tl.store(partial_sq_ptr + partial_index, sum_sq)
+
+
+@triton.jit
+def _group_norm_finalize_stats_kernel(
+    partial_sum_ptr,
+    partial_sq_ptr,
+    stats_ptr,
+    chunks_per_row,
+    group_size,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    sum_val = tl.zeros((), dtype=tl.float32)
+    sum_sq = tl.zeros((), dtype=tl.float32)
+    base = row * chunks_per_row
+    for off in range(0, chunks_per_row, BLOCK_SIZE):
+        idx = off + offsets
+        mask = idx < chunks_per_row
+        sum_val += tl.sum(
+            tl.load(partial_sum_ptr + base + idx, mask=mask, other=0.0), axis=0
+        )
+        sum_sq += tl.sum(
+            tl.load(partial_sq_ptr + base + idx, mask=mask, other=0.0), axis=0
+        )
+
+    inv_group = 1.0 / group_size
+    mean = sum_val * inv_group
+    var = sum_sq * inv_group - mean * mean
+    rstd = tl.rsqrt(var + eps)
+    tl.store(stats_ptr + row * 2, mean)
+    tl.store(stats_ptr + row * 2 + 1, rstd)
+
+
+@triton.jit
+def _group_norm_apply_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    stats_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+    weight_group_offset = group_id * channels_per_group
+
+    mean = tl.load(stats_ptr + row * 2)
+    rstd = tl.load(stats_ptr + row * 2 + 1)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        channel_offsets = weight_group_offset + idx // spatial_size
+        weight = tl.load(weight_ptr + channel_offsets, mask=mask, other=1.0).to(
+            tl.float32
+        )
+        bias = tl.load(bias_ptr + channel_offsets, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+@triton.jit
+def _group_norm_apply_scalar_affine_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    stats_ptr,
+    channels,
+    spatial_size,
+    num_groups,
+    channels_per_group,
+    group_size,
+    chunks_per_row,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCKS_PER_PROGRAM: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    chunk_id = tl.program_id(1).to(tl.int64)
+
+    batch_id = row // num_groups
+    group_id = row - batch_id * num_groups
+    chunk_start = chunk_id * BLOCK_SIZE * BLOCKS_PER_PROGRAM
+    group_base = batch_id * channels * spatial_size + group_id * group_size
+
+    channel_id = chunk_start // spatial_size
+    affine_offset = group_id * channels_per_group + channel_id
+    weight = tl.load(weight_ptr + affine_offset).to(tl.float32)
+    bias = tl.load(bias_ptr + affine_offset).to(tl.float32)
+
+    mean = tl.load(stats_ptr + row * 2)
+    rstd = tl.load(stats_ptr + row * 2 + 1)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    for block_id in range(BLOCKS_PER_PROGRAM):
+        idx = chunk_start + block_id * BLOCK_SIZE + offsets
+        mask = idx < group_size
+        x = tl.load(input_ptr + group_base + idx, mask=mask, other=0.0).to(tl.float32)
+        y = (x - mean) * rstd
+        y = y * weight + bias
+        y = y * tl.sigmoid(y)
+        tl.store(output_ptr + group_base + idx, y, mask=mask)
+
+
+def _group_norm_silu_native(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    return F.silu(F.group_norm(x, num_groups, weight=weight, bias=bias, eps=eps))
+
+
+def _can_use_triton_group_norm_silu(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+) -> bool:
+    return (
+        x.is_cuda
+        and not torch.is_grad_enabled()
+        and not x.requires_grad
+        and x.dtype in _SUPPORTED_DTYPES
+        and x.ndim in (2, 3, 4, 5)
+        and x.shape[1] % num_groups == 0
+        and weight.is_cuda
+        and bias.is_cuda
+        and weight.dtype == x.dtype
+        and bias.dtype == x.dtype
+        and weight.ndim == 1
+        and bias.ndim == 1
+        and weight.shape == bias.shape == (x.shape[1],)
+    )
+
+
+def _launch_one_pass(
+    x_contiguous: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    batch_size, channels = x_contiguous.shape[:2]
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = channels // num_groups
+    group_size = channels_per_group * spatial_size
+
+    x_flat = x_contiguous.reshape(batch_size, channels, spatial_size, 1)
+    y_flat = torch.empty_like(x_flat)
+    block_size = min(4096, triton.next_power_of_2(max(1, min(group_size, 4096))))
+
+    _group_norm_silu_contiguous_kernel[(num_groups, batch_size)](
+        x_flat,
+        weight,
+        bias,
+        y_flat,
+        channels,
+        spatial_size,
+        channels_per_group,
+        group_size,
+        eps,
+        BLOCK_SIZE=block_size,
+    )
+    return y_flat.reshape_as(x_contiguous)
+
+
+def _launch_chunked(
+    x_contiguous: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float,
+) -> torch.Tensor:
+    batch_size, channels = x_contiguous.shape[:2]
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = channels // num_groups
+    group_size = channels_per_group * spatial_size
+    rows = batch_size * num_groups
+    chunks_per_row = triton.cdiv(group_size, _CHUNK_SIZE)
+
+    x_flat = x_contiguous.reshape(-1)
+    y = torch.empty_like(x_contiguous)
+    y_flat = y.reshape(-1)
+    partial_sum = torch.empty(
+        (rows, chunks_per_row), device=x_contiguous.device, dtype=torch.float32
+    )
+    partial_sq = torch.empty_like(partial_sum)
+    stats = torch.empty((rows, 2), device=x_contiguous.device, dtype=torch.float32)
+
+    _group_norm_stats_kernel[(rows, chunks_per_row)](
+        x_flat,
+        partial_sum,
+        partial_sq,
+        channels,
+        spatial_size,
+        num_groups,
+        channels_per_group,
+        group_size,
+        chunks_per_row,
+        BLOCK_SIZE=_BLOCK_SIZE,
+        BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+        num_warps=8,
+        num_stages=3,
+    )
+
+    reduce_block = min(1024, triton.next_power_of_2(max(1, chunks_per_row)))
+    _group_norm_finalize_stats_kernel[(rows,)](
+        partial_sum,
+        partial_sq,
+        stats,
+        chunks_per_row,
+        group_size,
+        eps,
+        BLOCK_SIZE=reduce_block,
+        num_warps=4,
+        num_stages=2,
+    )
+
+    if spatial_size % _CHUNK_SIZE == 0 and chunks_per_row >= 64:
+        _group_norm_apply_scalar_affine_kernel[(rows, chunks_per_row)](
+            x_flat,
+            weight,
+            bias,
+            y_flat,
+            stats,
+            channels,
+            spatial_size,
+            num_groups,
+            channels_per_group,
+            group_size,
+            chunks_per_row,
+            BLOCK_SIZE=_BLOCK_SIZE,
+            BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+            num_warps=4,
+            num_stages=3,
+        )
+    else:
+        _group_norm_apply_kernel[(rows, chunks_per_row)](
+            x_flat,
+            weight,
+            bias,
+            y_flat,
+            stats,
+            channels,
+            spatial_size,
+            num_groups,
+            channels_per_group,
+            group_size,
+            chunks_per_row,
+            BLOCK_SIZE=_BLOCK_SIZE,
+            BLOCKS_PER_PROGRAM=_BLOCKS_PER_PROGRAM,
+            num_warps=8,
+            num_stages=3,
+        )
+    return y
+
+
+@register_custom_op(op_name="triton_group_norm_silu_cuda", out_shape="x")
+def _triton_group_norm_silu_cuda(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    if not _can_use_triton_group_norm_silu(x, weight, bias, num_groups):
+        return _group_norm_silu_native(x, weight, bias, num_groups, eps)
+
+    x_contiguous = x.contiguous()
+    spatial_size = math.prod(x_contiguous.shape[2:]) if x_contiguous.ndim > 2 else 1
+    channels_per_group = x_contiguous.shape[1] // num_groups
+    group_size = channels_per_group * spatial_size
+
+    with torch.cuda.device(x.device):
+        if group_size >= _LARGE_GROUP_THRESHOLD:
+            return _launch_chunked(x_contiguous, weight, bias, num_groups, eps)
+        return _launch_one_pass(x_contiguous, weight, bias, num_groups, eps)
+
+
+def triton_group_norm_silu(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    num_groups: int,
+    eps: float = 1e-5,
+) -> torch.Tensor:
+    return _triton_group_norm_silu_cuda(x, weight, bias, num_groups, eps)
+
+
+__all__ = ["triton_group_norm_silu"]

python/sglang/jit_kernel/diffusion/triton/ltx2_ada_values.py

@@ -0,0 +1,184 @@
+# Adapted from NVlabs/Sana sol-engine LTX2 Ada-value fusion.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _ltx2_ada_values9_kernel(
+    temb_ptr,
+    table_ptr,
+    out0_ptr,
+    out1_ptr,
+    out2_ptr,
+    out3_ptr,
+    out4_ptr,
+    out5_ptr,
+    out6_ptr,
+    out7_ptr,
+    out8_ptr,
+    rows: tl.constexpr,
+    hidden: tl.constexpr,
+    total_params: tl.constexpr,
+    table_stride_p: tl.constexpr,
+    table_stride_d: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < hidden
+    temb_row = temb_ptr + row * total_params * hidden
+    base = row * hidden + cols
+
+    table0 = tl.load(
+        table_ptr + 0 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb0 = tl.load(
+        temb_row + (0 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table1 = tl.load(
+        table_ptr + 1 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb1 = tl.load(
+        temb_row + (1 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table2 = tl.load(
+        table_ptr + 2 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb2 = tl.load(
+        temb_row + (2 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table3 = tl.load(
+        table_ptr + 3 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb3 = tl.load(
+        temb_row + (3 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table4 = tl.load(
+        table_ptr + 4 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb4 = tl.load(
+        temb_row + (4 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table5 = tl.load(
+        table_ptr + 5 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb5 = tl.load(
+        temb_row + (5 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table6 = tl.load(
+        table_ptr + 6 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb6 = tl.load(
+        temb_row + (6 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table7 = tl.load(
+        table_ptr + 7 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb7 = tl.load(
+        temb_row + (7 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    table8 = tl.load(
+        table_ptr + 8 * table_stride_p + cols * table_stride_d,
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+    temb8 = tl.load(
+        temb_row + (8 * hidden + cols),
+        mask=mask,
+        other=0.0,
+    ).to(tl.bfloat16)
+
+    tl.store(out0_ptr + base, (table0 + temb0).to(tl.bfloat16), mask=mask)
+    tl.store(out1_ptr + base, (table1 + temb1).to(tl.bfloat16), mask=mask)
+    tl.store(out2_ptr + base, (table2 + temb2).to(tl.bfloat16), mask=mask)
+    tl.store(out3_ptr + base, (table3 + temb3).to(tl.bfloat16), mask=mask)
+    tl.store(out4_ptr + base, (table4 + temb4).to(tl.bfloat16), mask=mask)
+    tl.store(out5_ptr + base, (table5 + temb5).to(tl.bfloat16), mask=mask)
+    tl.store(out6_ptr + base, (table6 + temb6).to(tl.bfloat16), mask=mask)
+    tl.store(out7_ptr + base, (table7 + temb7).to(tl.bfloat16), mask=mask)
+    tl.store(out8_ptr + base, (table8 + temb8).to(tl.bfloat16), mask=mask)
+
+
+def ltx2_ada_values9(
+    scale_shift_table: torch.Tensor,
+    timestep: torch.Tensor,
+) -> tuple[torch.Tensor, ...]:
+    if timestep.ndim != 3:
+        raise ValueError("timestep must have shape [B, S, 9 * D]")
+    if not timestep.is_cuda or timestep.dtype != torch.bfloat16:
+        raise ValueError("timestep must be a CUDA bfloat16 tensor")
+    if not timestep.is_contiguous():
+        raise ValueError("timestep must be contiguous")
+    if scale_shift_table.ndim != 2 or scale_shift_table.shape[0] != 9:
+        raise ValueError("scale_shift_table must have shape [9, D]")
+    if (
+        not scale_shift_table.is_cuda
+        or scale_shift_table.dtype not in (torch.bfloat16, torch.float32)
+        or scale_shift_table.stride(-1) != 1
+    ):
+        raise ValueError(
+            "scale_shift_table must be CUDA, bf16/fp32, last-dim contiguous"
+        )
+
+    total_params = int(scale_shift_table.shape[0])
+    hidden = int(scale_shift_table.shape[1])
+    if hidden <= 0 or timestep.shape[-1] != total_params * hidden:
+        raise ValueError("timestep last dim must equal 9 * hidden")
+    if hidden % 256 != 0 or hidden > 8192:
+        raise ValueError("hidden size is outside the supported LTX2 fast-path range")
+
+    batch, seq, _ = timestep.shape
+    rows = int(batch * seq)
+    outs = tuple(
+        torch.empty((batch, seq, hidden), device=timestep.device, dtype=timestep.dtype)
+        for _ in range(9)
+    )
+    _ltx2_ada_values9_kernel[(rows,)](
+        timestep,
+        scale_shift_table,
+        *outs,
+        rows,
+        hidden,
+        total_params,
+        scale_shift_table.stride(0),
+        scale_shift_table.stride(1),
+        BLOCK_N=triton.next_power_of_2(hidden),
+        num_warps=4 if hidden >= 4096 else 8,
+    )
+    return outs

python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py

@@ -0,0 +1,102 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _ltx2_split_rotary_kernel(
+    out_ptr,
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    seq_len: tl.constexpr,
+    num_heads: tl.constexpr,
+    head_dim: tl.constexpr,
+    half_dim: tl.constexpr,
+    stride_cos_b: tl.constexpr,
+    stride_cos_h: tl.constexpr,
+    stride_cos_t: tl.constexpr,
+    stride_sin_b: tl.constexpr,
+    stride_sin_h: tl.constexpr,
+    stride_sin_t: tl.constexpr,
+    BLOCK_HEADS: tl.constexpr,
+    BLOCK_HALF: tl.constexpr,
+):
+    pid_bt = tl.program_id(0)
+    head_block = tl.program_id(1)
+    batch = pid_bt // seq_len
+    token = pid_bt - batch * seq_len
+    heads = head_block * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)
+    offsets = tl.arange(0, BLOCK_HALF)
+    mask = (heads[:, None] < num_heads) & (offsets[None, :] < half_dim)
+
+    x_base = ((batch * seq_len + token) * num_heads + heads[:, None]) * head_dim
+    cos_base = (
+        batch * stride_cos_b + heads[:, None] * stride_cos_h + token * stride_cos_t
+    )
+    sin_base = (
+        batch * stride_sin_b + heads[:, None] * stride_sin_h + token * stride_sin_t
+    )
+
+    x_first = tl.load(x_ptr + x_base + offsets[None, :], mask=mask, other=0.0)
+    x_second = tl.load(
+        x_ptr + x_base + half_dim + offsets[None, :], mask=mask, other=0.0
+    )
+    cos = tl.load(cos_ptr + cos_base + offsets[None, :], mask=mask, other=0.0)
+    sin = tl.load(sin_ptr + sin_base + offsets[None, :], mask=mask, other=0.0)
+
+    # Match the original PyTorch order: x * cos is written as BF16 first, then
+    # addcmul_ computes the sine product in FP32 before the final BF16 store.
+    out_first = (x_first * cos).to(tl.bfloat16).to(tl.float32) + (
+        -x_second.to(tl.float32) * sin.to(tl.float32)
+    )
+    out_second = (x_second * cos).to(tl.bfloat16).to(tl.float32) + (
+        x_first.to(tl.float32) * sin.to(tl.float32)
+    )
+
+    tl.store(out_ptr + x_base + offsets[None, :], out_first, mask=mask)
+    tl.store(out_ptr + x_base + half_dim + offsets[None, :], out_second, mask=mask)
+
+
+def apply_ltx2_split_rotary_emb(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> torch.Tensor:
+    batch, seq_len, inner_dim = x.shape
+    cos_batch, num_heads, cos_seq_len, half_dim = cos.shape
+    head_dim = half_dim * 2
+    if (
+        cos_batch != batch
+        or cos_seq_len != seq_len
+        or inner_dim != num_heads * head_dim
+        or sin.shape != cos.shape
+    ):
+        raise ValueError(
+            "LTX2 split RoPE shape mismatch: "
+            f"x={tuple(x.shape)}, cos={tuple(cos.shape)}, sin={tuple(sin.shape)}"
+        )
+
+    out = torch.empty_like(x)
+    block_half = triton.next_power_of_2(half_dim)
+    block_heads = min(16, triton.next_power_of_2(num_heads))
+    num_warps = min(8, max(1, block_heads))
+    grid = (batch * seq_len, triton.cdiv(num_heads, block_heads))
+    _ltx2_split_rotary_kernel[grid](
+        out,
+        x,
+        cos,
+        sin,
+        seq_len,
+        num_heads,
+        head_dim,
+        half_dim,
+        cos.stride(0),
+        cos.stride(1),
+        cos.stride(2),
+        sin.stride(0),
+        sin.stride(1),
+        sin.stride(2),
+        BLOCK_HEADS=block_heads,
+        BLOCK_HALF=block_half,
+        num_warps=num_warps,
+    )
+    return out

python/sglang/jit_kernel/diffusion/triton/mps_fallback.py

@@ -0,0 +1,125 @@
+"""MPS (Apple Silicon) fallbacks for Triton diffusion kernels.
+
+Triton is not available on macOS / Metal, so these pure-PyTorch (and
+optionally MLX-accelerated) implementations replace the Triton kernels
+at import time when ``current_platform.is_mps()`` is True.
+
+MLX acceleration (opt-in via ``SGLANG_USE_MLX=1``):
+    Norm ops use ``mx.fast.rms_norm`` / ``mx.fast.layer_norm`` — single fused
+    Metal kernels that are 1.4x–2.9x faster than the multi-step PyTorch MPS
+    decomposition for medium-to-large tensors.
+"""
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from sglang.srt.utils.tensor_bridge import mlx_to_torch, torch_to_mlx, use_mlx
+
+from .torch_fallback import (
+    apply_rotary_embedding_native,
+    fuse_scale_shift_kernel_native,
+    norm_infer_native,
+    rms_norm_fn_native,
+    triton_one_pass_rms_norm_native,
+)
+
+_use_mlx = use_mlx()
+
+if _use_mlx:
+    import mlx.core as mx
+
+# use the common torch native version form torch_fallback
+fuse_scale_shift_kernel_native = fuse_scale_shift_kernel_native
+apply_rotary_embedding_native = apply_rotary_embedding_native
+norm_infer_native = norm_infer_native
+triton_one_pass_rms_norm_native = triton_one_pass_rms_norm_native
+rms_norm_fn_native = rms_norm_fn_native
+
+# MLX-accelerated norm ops (1.4x–2.9x faster than torch native on MPS)
+# Uses mx.fast.rms_norm / mx.fast.layer_norm — single fused Metal kernels
+# instead of 7+ separate PyTorch MPS kernel launches.
+
+if _use_mlx:
+
+    def norm_infer_native(  # noqa: F811
+        x: Tensor,
+        weight: Optional[Tensor],
+        bias: Optional[Tensor],
+        eps: float,
+        is_rms_norm: bool = False,
+        out: Optional[Tensor] = None,
+    ) -> Tensor:
+        """MLX-accelerated norm_infer (layer norm / rms norm inference)."""
+        device = x.device
+        orig_dtype = x.dtype
+        x_mx = torch_to_mlx(x)
+        if is_rms_norm:
+            w_mx = (
+                torch_to_mlx(weight) if weight is not None else mx.ones(x_mx.shape[-1])
+            )
+            result_mx = mx.fast.rms_norm(x_mx, w_mx, eps)
+        else:
+            w_mx = torch_to_mlx(weight) if weight is not None else None
+            b_mx = torch_to_mlx(bias) if bias is not None else None
+            result_mx = mx.fast.layer_norm(x_mx, w_mx, b_mx, eps)
+        result = mlx_to_torch(result_mx, device).to(orig_dtype)
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+
+    def triton_one_pass_rms_norm_native(  # noqa: F811
+        x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
+    ) -> torch.Tensor:
+        """MLX-accelerated triton_one_pass_rms_norm."""
+        device = x.device
+        orig_dtype = x.dtype
+        x_mx = torch_to_mlx(x)
+        w_mx = torch_to_mlx(w)
+        result_mx = mx.fast.rms_norm(x_mx, w_mx, eps)
+        return mlx_to_torch(result_mx, device).to(orig_dtype)
+
+    def rms_norm_fn_native(  # noqa: F811
+        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,
+    ):
+        """MLX-accelerated rms_norm_fn (inference only, no dropout/x1 support)."""
+        device = x.device
+        orig_dtype = x.dtype
+        if residual is not None:
+            x = x.float() + residual.float()
+            residual_out_val = x.to(torch.float32 if residual_in_fp32 else orig_dtype)
+        else:
+            residual_out_val = None
+        if weight is not None and zero_centered_weight:
+            w = weight.float() + 1.0
+        else:
+            w = weight
+        x_mx = torch_to_mlx(x)
+        w_mx = torch_to_mlx(w) if w is not None else mx.ones(x_mx.shape[-1])
+        result_mx = mx.fast.rms_norm(x_mx, w_mx, eps)
+        x_hat = mlx_to_torch(result_mx, device)
+        if bias is not None:
+            x_hat = x_hat + bias.to(x_hat.device, x_hat.dtype)
+        final_dtype = out_dtype if out_dtype is not None else orig_dtype
+        y = x_hat.to(final_dtype)
+        if residual is not None and residual_out_val is not None:
+            return y, residual_out_val
+        return y

python/sglang/jit_kernel/diffusion/triton/norm.py

@@ -0,0 +1,660 @@
+from typing import Optional, Tuple
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+from torch import Tensor
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+# RMSNorm-fp32
+def maybe_contiguous_lastdim(x):
+    return x.contiguous() if x is not None and x.stride(-1) != 1 else x
+
+
+def maybe_contiguous(x):
+    return x.contiguous() if x is not None else None
+
+
+def triton_autotune_configs():
+    # Return configs with a valid warp count for the current device
+    # 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
+    warp_size = getattr(
+        torch.get_device_module().get_device_properties(
+            torch.get_device_module().current_device()
+        ),
+        "warp_size",
+        32,
+    )
+    if warp_size is None:
+        warp_size = 32
+    # Autotune for warp counts which are powers of 2 and do not exceed thread per block limit
+    return [
+        triton.Config({}, num_warps=warp_count)
+        for warp_count in [1, 2, 4, 8, 16, 32]
+        if warp_count * warp_size <= max_threads_per_block
+    ]
+
+
+# Copied from flash-attn
+@triton.autotune(
+    configs=triton_autotune_configs(),
+    key=[
+        "N",
+        "HAS_RESIDUAL",
+        "STORE_RESIDUAL_OUT",
+        "IS_RMS_NORM",
+        "HAS_BIAS",
+        "HAS_WEIGHT",
+        "HAS_X1",
+        "HAS_W1",
+        "HAS_B1",
+    ],
+)
+# torch compile doesn't like triton.heuristics, so we set these manually when calling the kernel
+# @triton.heuristics({"HAS_BIAS": lambda args: args["B"] is not None})
+# @triton.heuristics({"HAS_RESIDUAL": lambda args: args["RESIDUAL"] is not None})
+# @triton.heuristics({"HAS_X1": lambda args: args["X1"] is not None})
+# @triton.heuristics({"HAS_W1": lambda args: args["W1"] is not None})
+# @triton.heuristics({"HAS_B1": lambda args: args["B1"] is not None})
+@triton.jit
+def _layer_norm_fwd_1pass_kernel(
+    X,  # pointer to the input
+    Y,  # pointer to the output
+    W,  # pointer to the weights
+    B,  # pointer to the biases
+    RESIDUAL,  # pointer to the residual
+    X1,
+    W1,
+    B1,
+    Y1,
+    RESIDUAL_OUT,  # pointer to the residual
+    ROWSCALE,
+    SEEDS,  # Dropout seeds for each row
+    DROPOUT_MASK,
+    DROPOUT_MASK1,
+    Mean,  # pointer to the mean
+    Rstd,  # pointer to the 1/std
+    stride_x_row,  # how much to increase the pointer when moving by 1 row
+    stride_y_row,
+    stride_res_row,
+    stride_res_out_row,
+    stride_x1_row,
+    stride_y1_row,
+    M,  # number of rows in X
+    N,  # number of columns in X
+    eps,  # epsilon to avoid division by zero
+    dropout_p,  # Dropout probability
+    zero_centered_weight,  # If true, add 1.0 to the weight
+    IS_RMS_NORM: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    HAS_RESIDUAL: tl.constexpr,
+    STORE_RESIDUAL_OUT: tl.constexpr,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    HAS_DROPOUT: tl.constexpr,
+    STORE_DROPOUT_MASK: tl.constexpr,
+    HAS_ROWSCALE: tl.constexpr,
+    HAS_X1: tl.constexpr,
+    HAS_W1: tl.constexpr,
+    HAS_B1: tl.constexpr,
+):
+    # Map the program id to the row of X and Y it should compute.
+    row = tl.program_id(0)
+    X += row * stride_x_row
+    Y += row * stride_y_row
+    if HAS_RESIDUAL:
+        RESIDUAL += row * stride_res_row
+    if STORE_RESIDUAL_OUT:
+        RESIDUAL_OUT += row * stride_res_out_row
+    if HAS_X1:
+        X1 += row * stride_x1_row
+    if HAS_W1:
+        Y1 += row * stride_y1_row
+    # Compute mean and variance
+    cols = tl.arange(0, BLOCK_N)
+    x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
+    if HAS_ROWSCALE:
+        rowscale = tl.load(ROWSCALE + row).to(tl.float32)
+        x *= rowscale
+    if HAS_DROPOUT:
+        # Compute dropout mask
+        # 7 rounds is good enough, and reduces register pressure
+        keep_mask = (
+            tl.rand(tl.load(SEEDS + row).to(tl.uint32), cols, n_rounds=7) > dropout_p
+        )
+        x = tl.where(keep_mask, x / (1.0 - dropout_p), 0.0)
+        if STORE_DROPOUT_MASK:
+            tl.store(DROPOUT_MASK + row * N + cols, keep_mask, mask=cols < N)
+    if HAS_X1:
+        x1 = tl.load(X1 + cols, mask=cols < N, other=0.0).to(tl.float32)
+        if HAS_ROWSCALE:
+            rowscale = tl.load(ROWSCALE + M + row).to(tl.float32)
+            x1 *= rowscale
+        if HAS_DROPOUT:
+            # Compute dropout mask
+            # 7 rounds is good enough, and reduces register pressure
+            keep_mask = (
+                tl.rand(tl.load(SEEDS + M + row).to(tl.uint32), cols, n_rounds=7)
+                > dropout_p
+            )
+            x1 = tl.where(keep_mask, x1 / (1.0 - dropout_p), 0.0)
+            if STORE_DROPOUT_MASK:
+                tl.store(DROPOUT_MASK1 + row * N + cols, keep_mask, mask=cols < N)
+        x += x1
+    if HAS_RESIDUAL:
+        residual = tl.load(RESIDUAL + cols, mask=cols < N, other=0.0).to(tl.float32)
+        x += residual
+    if STORE_RESIDUAL_OUT:
+        tl.store(RESIDUAL_OUT + cols, x, mask=cols < N)
+    if not IS_RMS_NORM:
+        mean = tl.sum(x, axis=0) / N
+        tl.store(Mean + row, mean)
+        xbar = tl.where(cols < N, x - mean, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    else:
+        xbar = tl.where(cols < N, x, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    rstd = 1 / tl.sqrt(var + eps)
+    tl.store(Rstd + row, rstd)
+    # Normalize and apply linear transformation
+    mask = cols < N
+    if HAS_WEIGHT:
+        w = tl.load(W + cols, mask=mask).to(tl.float32)
+        if zero_centered_weight:
+            w += 1.0
+    if HAS_BIAS:
+        b = tl.load(B + cols, mask=mask).to(tl.float32)
+    x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
+    if HAS_WEIGHT:
+        y = x_hat * w + b if HAS_BIAS else x_hat * w
+    else:
+        y = x_hat + b if HAS_BIAS else x_hat
+    # Write output
+    tl.store(Y + cols, y, mask=mask)
+    if HAS_W1:
+        w1 = tl.load(W1 + cols, mask=mask).to(tl.float32)
+        if zero_centered_weight:
+            w1 += 1.0
+        if HAS_B1:
+            b1 = tl.load(B1 + cols, mask=mask).to(tl.float32)
+        y1 = x_hat * w1 + b1 if HAS_B1 else x_hat * w1
+        tl.store(Y1 + cols, y1, mask=mask)
+
+
+def _layer_norm_fwd(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    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,
+    out_dtype: Optional[torch.dtype] = None,
+    residual_dtype: Optional[torch.dtype] = None,
+    zero_centered_weight: bool = False,
+    is_rms_norm: bool = False,
+    return_dropout_mask: bool = False,
+    out: Optional[Tensor] = None,
+    residual_out: Optional[Tensor] = None,
+) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor):
+    # 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:
+        residual_dtype = residual.dtype
+    if residual_out is None and (
+        residual is not None
+        or (residual_dtype is not None and residual_dtype != x.dtype)
+        or dropout_p > 0.0
+        or rowscale is not None
+        or x1 is not None
+    ):
+        residual_out = torch.empty_like(
+            x, dtype=residual_dtype if residual_dtype is not None else x.dtype
+        )
+    else:
+        residual_out = None
+    y1, mean, rstd, seeds, dropout_mask, dropout_mask1 = _layer_norm_fwd_impl(
+        x,
+        weight,
+        bias,
+        eps,
+        out,
+        residual=residual,
+        x1=x1,
+        weight1=weight1,
+        bias1=bias1,
+        dropout_p=dropout_p,
+        rowscale=rowscale,
+        zero_centered_weight=zero_centered_weight,
+        is_rms_norm=is_rms_norm,
+        return_dropout_mask=return_dropout_mask,
+        residual_out=residual_out,
+    )
+    # residual_out is None if residual is None and residual_dtype == input_dtype and dropout_p == 0.0
+    if residual_out is None:
+        residual_out = x
+    return out, y1, mean, rstd, residual_out, seeds, dropout_mask, dropout_mask1
+
+
+@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: 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,
+    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,
+) -> None:
+    M, N = x.shape
+    assert x.stride(-1) == 1
+    if residual is not None:
+        assert residual.stride(-1) == 1
+        assert residual.shape == (M, N)
+    if weight is not None:
+        assert weight.shape == (N,)
+        assert weight.stride(-1) == 1
+    if bias is not None:
+        assert bias.stride(-1) == 1
+        assert bias.shape == (N,)
+    if x1 is not None:
+        assert x1.shape == x.shape
+        assert rowscale is None
+        assert x1.stride(-1) == 1
+    if weight1 is not None:
+        assert weight1.shape == (N,)
+        assert weight1.stride(-1) == 1
+    if bias1 is not None:
+        assert bias1.shape == (N,)
+        assert bias1.stride(-1) == 1
+    if rowscale is not None:
+        assert rowscale.is_contiguous()
+        assert rowscale.shape == (M,)
+    assert out.shape == x.shape
+    assert out.stride(-1) == 1
+    if residual_out is not None:
+        assert residual_out.shape == x.shape
+        assert residual_out.stride(-1) == 1
+    if y1 is not None:
+        assert y1.shape == x.shape
+        assert y1.stride(-1) == 1
+    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):
+        _layer_norm_fwd_1pass_kernel[(M,)](
+            x,
+            out,
+            weight if weight is not None else x,  # unused when HAS_WEIGHT == False
+            bias,
+            residual,
+            x1,
+            weight1,
+            bias1,
+            y1,
+            residual_out,
+            rowscale,
+            seeds,
+            dropout_mask,
+            dropout_mask1,
+            mean,
+            rstd,
+            x.stride(0),
+            out.stride(0),
+            residual.stride(0) if residual is not None else 0,
+            residual_out.stride(0) if residual_out is not None else 0,
+            x1.stride(0) if x1 is not None else 0,
+            y1.stride(0) if y1 is not None else 0,
+            M,
+            N,
+            eps,
+            dropout_p,
+            # Passing bool make torch inductor very unhappy since it then tries to compare to int_max
+            int(zero_centered_weight),
+            is_rms_norm,
+            BLOCK_N,
+            residual is not None,
+            residual_out is not None,
+            weight is not None,
+            bias is not None,
+            dropout_p > 0.0,
+            dropout_mask is not None,
+            rowscale is not None,
+            HAS_X1=x1 is not None,
+            HAS_W1=weight1 is not None,
+            HAS_B1=bias1 is not None,
+        )
+    return None
+
+
+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,
+        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
+
+
+def _norm_forward(
+    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
+
+
+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,
+        residual,
+        x1,
+        weight1,
+        bias1,
+        eps,
+        dropout_p,
+        rowscale,
+        prenorm,
+        residual_in_fp32,
+        zero_centered_weight,
+        True,
+        return_dropout_mask,
+        out_dtype,
+        out,
+        residual_out,
+    )
+
+
+@triton.jit
+def _norm_infer_kernel(
+    X,
+    Y,
+    W,
+    B,
+    stride_x_row,
+    stride_y_row,
+    M,
+    N,
+    eps,
+    IS_RMS_NORM: tl.constexpr,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    X += row * stride_x_row
+    Y += row * stride_y_row
+    if HAS_WEIGHT:
+        W += 0
+    if HAS_BIAS:
+        B += 0
+    cols = tl.arange(0, BLOCK_N)
+    x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32)
+    if not IS_RMS_NORM:
+        mean = tl.sum(x, axis=0) / N
+        xbar = tl.where(cols < N, x - mean, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    else:
+        xbar = tl.where(cols < N, x, 0.0)
+        var = tl.sum(xbar * xbar, axis=0) / N
+    rstd = 1 / tl.sqrt(var + eps)
+    x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
+    if HAS_WEIGHT:
+        w = tl.load(W + cols, mask=cols < N, other=1.0).to(tl.float32)
+        y = x_hat * w
+    else:
+        y = x_hat
+    if HAS_BIAS:
+        b = tl.load(B + cols, mask=cols < N, other=0.0).to(tl.float32)
+        y += b
+    tl.store(Y + cols, y, mask=cols < N)
+
+
+def norm_infer(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    is_rms_norm: bool = False,
+    out: Optional[Tensor] = None,
+):
+    M, N = x.shape
+    x = x.contiguous()
+    if weight is not None:
+        assert weight.shape == (N,)
+        assert weight.stride(-1) == 1
+    if bias is not None:
+        assert bias.shape == (N,)
+        assert bias.stride(-1) == 1
+    if out is None:
+        out = torch.empty_like(x)
+    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.")
+    num_warps = min(max(BLOCK_N // 256, 1), 8)
+    _norm_infer_kernel[(M,)](
+        x,
+        out,
+        weight if weight is not None else x,  # dummy when HAS_WEIGHT=False
+        bias if bias is not None else x,  # dummy when HAS_BIAS=False
+        x.stride(0),
+        out.stride(0),
+        M,
+        N,
+        eps,
+        IS_RMS_NORM=is_rms_norm,
+        HAS_WEIGHT=weight is not None,
+        HAS_BIAS=bias is not None,
+        BLOCK_N=BLOCK_N,
+        num_warps=num_warps,
+    )
+    return out
+
+
+if current_platform.is_mps():
+    from .mps_fallback import norm_infer_native, rms_norm_fn_native
+
+    norm_infer = norm_infer_native
+    rms_norm_fn = rms_norm_fn_native
+
+if current_platform.is_cpu():
+    from .torch_fallback import norm_infer_native, rms_norm_fn_native
+
+    norm_infer = norm_infer_native
+    rms_norm_fn = rms_norm_fn_native

python/sglang/jit_kernel/diffusion/triton/npu_fallback.py

@@ -0,0 +1,50 @@
+import torch
+import torch_npu
+
+NPU_ROTARY_MUL_MAX_NUM_HEADS = 1000
+NPU_ROTARY_MUL_MAX_HEAD_SIZE = 896
+
+
+# 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:
+    if interleaved and cos.shape[-1] == x.shape[-1]:
+        cos = cos[..., ::2]
+        sin = sin[..., ::2]
+    cos = cos.unsqueeze(-2).to(x.dtype)
+    sin = sin.unsqueeze(-2).to(x.dtype)
+
+    if (
+        cos.dim() == 3
+        and x.dim() == 3
+        and x.shape[1] < NPU_ROTARY_MUL_MAX_NUM_HEADS
+        and x.shape[2] < NPU_ROTARY_MUL_MAX_HEAD_SIZE
+        and not interleaved
+    ):
+        if cos.size(-1) * 2 == x.size(-1):
+            cos = torch.cat([cos, cos], dim=-1)
+            sin = torch.cat([sin, sin], dim=-1)
+        cos = cos.unsqueeze(0)
+        sin = sin.unsqueeze(0)
+        x = x.unsqueeze(0)
+        x_embed = torch_npu.npu_rotary_mul(x, cos, sin)
+        x_embed = x_embed.squeeze(0)
+        return x_embed
+
+    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,83 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+# 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)
+
+
+@register_custom_op(op_name="triton_one_pass_rms_norm_cuda", out_shape="x")
+def _triton_one_pass_rms_norm_cuda(
+    x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    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):
+        _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
+
+
+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)
+
+
+if current_platform.is_mps():
+    from .mps_fallback import triton_one_pass_rms_norm_native
+
+    @debug_kernel_api
+    def triton_one_pass_rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6):
+        return triton_one_pass_rms_norm_native(x, w, eps)
+
+
+if current_platform.is_cpu():
+    from .torch_fallback import triton_one_pass_rms_norm_native
+
+    triton_one_pass_rms_norm = triton_one_pass_rms_norm_native

python/sglang/jit_kernel/diffusion/triton/rotary.py

@@ -0,0 +1,141 @@
+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_HEADS": 1, "BLOCK_HS_HALF": 32}, num_warps=2),
+        triton.Config({"BLOCK_HEADS": 2, "BLOCK_HS_HALF": 32}, num_warps=2),
+        triton.Config({"BLOCK_HEADS": 4, "BLOCK_HS_HALF": 32}, num_warps=4),
+        triton.Config({"BLOCK_HEADS": 4, "BLOCK_HS_HALF": 64}, num_warps=4),
+        triton.Config({"BLOCK_HEADS": 8, "BLOCK_HS_HALF": 64}, num_warps=8),
+    ],
+    key=["num_heads", "head_size"],
+)
+@triton.jit
+def _rotary_embedding_kernel(
+    output_ptr,
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    num_heads,
+    head_size,
+    num_tokens,
+    stride_out_bt,
+    stride_out_head,
+    stride_x_bt,
+    stride_x_head,
+    stride_cos_row,
+    stride_sin_row,
+    BLOCK_HEADS: tl.constexpr,
+    BLOCK_HS_HALF: tl.constexpr,
+):
+    bt_idx = tl.program_id(0)
+    head_block_idx = tl.program_id(1)
+    token_idx = bt_idx % num_tokens
+
+    cos_row_ptr = cos_ptr + token_idx * stride_cos_row
+    sin_row_ptr = sin_ptr + token_idx * stride_sin_row
+    head_offsets = head_block_idx * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)
+    head_mask = head_offsets < num_heads
+
+    head_size_half = head_size // 2
+    x_row_ptrs = x_ptr + bt_idx * stride_x_bt + head_offsets[:, None] * stride_x_head
+    output_row_ptrs = (
+        output_ptr + bt_idx * stride_out_bt + head_offsets[:, None] * stride_out_head
+    )
+
+    for block_start in range(0, head_size_half, BLOCK_HS_HALF):
+        offsets_half = block_start + tl.arange(0, BLOCK_HS_HALF)
+        half_mask = offsets_half < head_size_half
+        mask = head_mask[:, None] & half_mask[None, :]
+
+        cos_vals = tl.load(cos_row_ptr + offsets_half, mask=half_mask, other=0.0)
+        sin_vals = tl.load(sin_row_ptr + offsets_half, mask=half_mask, other=0.0)
+
+        offsets_x1 = 2 * offsets_half
+        offsets_x2 = 2 * offsets_half + 1
+
+        x1_vals = tl.load(x_row_ptrs + offsets_x1[None, :], mask=mask, other=0.0)
+        x2_vals = tl.load(x_row_ptrs + offsets_x2[None, :], 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)[None, :]
+        sin_fp32 = sin_vals.to(tl.float32)[None, :]
+        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_ptrs + offsets_x1[None, :],
+            o1_vals.to(x1_vals.dtype),
+            mask=mask,
+        )
+        tl.store(
+            output_row_ptrs + offsets_x2[None, :],
+            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(bsz * num_tokens, num_heads, head_size)
+    output_reshaped = output.view(bsz * num_tokens, num_heads, head_size)
+
+    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[
+        lambda META: (bsz * num_tokens, triton.cdiv(num_heads, META["BLOCK_HEADS"]))
+    ](
+        output_reshaped,
+        x_reshaped,
+        cos,
+        sin,
+        num_heads,
+        head_size,
+        num_tokens,
+        output_reshaped.stride(0),
+        output_reshaped.stride(1),
+        x_reshaped.stride(0),
+        x_reshaped.stride(1),
+        cos.stride(0),
+        sin.stride(0),
+    )
+
+    return output
+
+
+if current_platform.is_npu():
+    from .npu_fallback import apply_rotary_embedding_native
+
+    apply_rotary_embedding = apply_rotary_embedding_native
+
+if current_platform.is_mps():
+    from .mps_fallback import apply_rotary_embedding_native
+
+    apply_rotary_embedding = apply_rotary_embedding_native
+
+if current_platform.is_cpu():
+    from .torch_fallback import apply_rotary_embedding_native
+
+    apply_rotary_embedding = apply_rotary_embedding_native

python/sglang/jit_kernel/diffusion/triton/sana_wm_gdn.py

@@ -0,0 +1,241 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Inference-side helpers for the bidirectional fused GDN path.
+
+Precision knob: env var ``FUSED_GDN_PRECISION`` or ``PRECISION_OVERRIDE``:
+  0=IEEE fp32 dots, 1=TF32, 2=bf16 TC + fp32 state [default], 3=bf16 TC + bf16 state.
+"""
+
+# ruff: noqa: E501
+
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+# =====================================================================
+#  GPU-adaptive kernel config
+# =====================================================================
+
+
+def _get_kernel_config() -> dict:
+    """Return optimal kernel parameters for the current GPU.
+
+    STATE_FP32 (fp32 state_prev) needs ~128KB SRAM (H100 228KB), vs ~96KB for
+    bf16 state_prev (fits GB10's 101KB).
+    """
+    if not torch.cuda.is_available():
+        return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 4, "STATE_FP32": False}
+    smem = torch.cuda.get_device_properties(0).shared_memory_per_multiprocessor
+    state_fp32 = smem >= 150 * 1024  # H100 (228KB) yes, GB10 (101KB) no
+    return {"BLOCK_S": 64, "num_stages": 1, "num_warps": 8, "STATE_FP32": state_fp32}
+
+
+_KCFG = None
+
+
+def _kcfg():
+    global _KCFG
+    if _KCFG is None:
+        _KCFG = _get_kernel_config()
+    return _KCFG
+
+
+# precision=0 → IEEE fp32 dots + fp32 state  (DOT_PRECISION=2, STATE_FP32=1)
+# precision=1 → TF32  dots   + fp32 state    (DOT_PRECISION=1, STATE_FP32=1)
+# precision=2 → bf16  dots   + fp32 state    (DOT_PRECISION=0, STATE_FP32=1) [default]
+# precision=3 → bf16  dots   + bf16 state    (DOT_PRECISION=0, STATE_FP32=0)
+def _precision_params(precision: int) -> tuple:
+    if precision == 0:
+        return 2, True
+    elif precision == 1:
+        return 1, True
+    elif precision == 3:
+        return 0, False
+    else:  # default
+        return 0, True
+
+
+_env_prec = os.environ.get("FUSED_GDN_PRECISION", None)
+PRECISION_OVERRIDE: int | None = int(_env_prec) if _env_prec is not None else None
+
+
+def _resolve_launch_config() -> tuple:
+    """Returns (prec, dot_prec, state_fp32, num_warps).
+
+    Uses ``PRECISION_OVERRIDE`` when set, else ``_kcfg()`` (per-GPU SRAM).
+    num_warps clamped to 4 when dots run on fp32 operands (more registers).
+    """
+    cfg = _kcfg()
+    prec = PRECISION_OVERRIDE if PRECISION_OVERRIDE is not None else 2
+    dot_prec, state_fp32 = _precision_params(prec)
+    if PRECISION_OVERRIDE is None:
+        state_fp32 = cfg["STATE_FP32"]
+    nw = cfg["num_warps"]
+    if dot_prec >= 1:
+        nw = min(nw, 4)
+    return prec, dot_prec, state_fp32, nw
+
+
+def prepare_rope_tables(
+    rotary_emb, N: int, D: int, device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Complex rotary_emb `(1, 1, N, D//2)` → expanded (N, D) cos/sin tables.
+
+    Encodes the interleaved-pair rotation
+        y[2i]   = x[2i]*cos[i] - x[2i+1]*sin[i]
+        y[2i+1] = x[2i]*sin[i] + x[2i+1]*cos[i]
+    as  y[d] = x[d]*cos_exp[d] + x[d^1]*sin_exp[d]
+    where sin_exp[2i] = -sin[i], sin_exp[2i+1] = +sin[i].
+
+    Returns (cos_exp, sin_exp) both (N, D) float32, contiguous.
+    """
+    if rotary_emb is None:
+        return (
+            torch.ones(N, D, device=device, dtype=torch.float32),
+            torch.zeros(N, D, device=device, dtype=torch.float32),
+        )
+    freqs = rotary_emb.squeeze(0).squeeze(0)  # (N, D//2) complex
+    cos_half = freqs.real.float()
+    sin_half = freqs.imag.float()
+    rope_cos = cos_half.repeat_interleave(2, dim=-1)
+    rope_sin = torch.stack([-sin_half, sin_half], dim=-1).reshape(N, D)
+    return rope_cos.contiguous(), rope_sin.contiguous()
+
+
+# =====================================================================
+#  Fused single-pass Q+K inverse-RMS Triton kernel
+# =====================================================================
+# Single Triton launch that reads each `(b, n)` row of `qkv` once and emits
+# both `q_inv_rms[b, n]` and `k_inv_rms[b, n]`. Replaces two separate PyTorch
+# scans (cast→square→sum→rsqrt) over `qkv[:, :, 0]` and `qkv[:, :, 1]`.
+#
+# Layout assumed: `qkv` is (B, N, 3, H, D) contiguous, so the C = H*D channels
+# for a given (b, n, qkv_idx) live in a contiguous memory span.
+
+
+@triton.jit
+def _fused_qk_inv_rms_kernel(
+    qkv_ptr,  # *T_in     (B, N, 3, H, D), contiguous
+    q_inv_rms_ptr,  # *float32  (B, N)
+    k_inv_rms_ptr,  # *float32  (B, N)
+    N: tl.constexpr,
+    C: tl.constexpr,  # H * D
+    eps,
+    BLOCK_C: tl.constexpr,
+):
+    bn_id = tl.program_id(0)
+    qkv_row_stride = 3 * C
+    row_base = bn_id * qkv_row_stride
+    q_base = row_base
+    k_base = row_base + C
+
+    offs = tl.arange(0, BLOCK_C)
+    mask = offs < C
+
+    q_vals = tl.load(qkv_ptr + q_base + offs, mask=mask, other=0.0).to(tl.float32)
+    k_vals = tl.load(qkv_ptr + k_base + offs, mask=mask, other=0.0).to(tl.float32)
+
+    q_sq = tl.sum(q_vals * q_vals, axis=0)
+    k_sq = tl.sum(k_vals * k_vals, axis=0)
+
+    inv_c = 1.0 / C
+    q_inv = tl.rsqrt(q_sq * inv_c + eps)
+    k_inv = tl.rsqrt(k_sq * inv_c + eps)
+
+    tl.store(q_inv_rms_ptr + bn_id, q_inv)
+    tl.store(k_inv_rms_ptr + bn_id, k_inv)
+
+
+def fused_qk_inv_rms(
+    qkv: torch.Tensor,
+    eps: float = 1e-5,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Single-pass Triton fused Q+K inverse-RMS.
+
+    Replaces two separate PyTorch RMS scans with one launch that reads each
+    ``(b, n)`` row of ``qkv`` exactly once.
+    qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32.
+    """
+    assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)"
+    assert (
+        qkv.dim() == 5 and qkv.shape[2] == 3
+    ), f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}"
+    B, N, _, H, D = qkv.shape
+    C = H * D
+    q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device)
+    k_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device)
+    BLOCK_C = triton.next_power_of_2(C)
+    _fused_qk_inv_rms_kernel[(B * N,)](
+        qkv,
+        q_inv_rms,
+        k_inv_rms,
+        N=N,
+        C=C,
+        eps=eps,
+        BLOCK_C=BLOCK_C,
+    )
+    return q_inv_rms, k_inv_rms
+
+
+# =====================================================================
+#  Bidirectional GDN entry point (delegates to chunkwise)
+# =====================================================================
+
+
+def fused_bigdn_func(
+    qkv: torch.Tensor,  # (B, N, 3, H, D)
+    q_inv_rms: torch.Tensor,  # (B, N) float32
+    k_inv_rms: torch.Tensor,  # (B, N) float32
+    q_norm_weight: torch.Tensor,  # (C,) float32
+    k_norm_weight: torch.Tensor,  # (C,) float32
+    rope_cos: torch.Tensor,  # (N, D) float32
+    rope_sin: torch.Tensor,  # (N, D) float32
+    beta: torch.Tensor,  # (B, H, F, S)
+    decay: torch.Tensor,  # (B, H, F)
+    F: int,
+    S: int,
+    k_scale: float,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Bidirectional fused GDN. Returns ``(B, N, H, D)``.
+
+    Thin entry point kept for call-site stability; delegates to
+    :func:`fused_bigdn_bidi_chunkwise` from ``sana_wm_gdn_chunkwise``.
+    """
+    from sglang.jit_kernel.diffusion.triton.sana_wm_gdn_chunkwise import (
+        fused_bigdn_bidi_chunkwise,
+    )
+
+    return fused_bigdn_bidi_chunkwise(
+        qkv,
+        q_inv_rms,
+        k_inv_rms,
+        q_norm_weight,
+        k_norm_weight,
+        rope_cos,
+        rope_sin,
+        beta,
+        decay,
+        F=F,
+        S=S,
+        k_scale=k_scale,
+        eps=eps,
+    )

python/sglang/jit_kernel/diffusion/triton/scale_shift.py

@@ -0,0 +1,691 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+from sglang.multimodal_gen.runtime.platforms import current_platform
+
+
+@triton.jit
+def _fused_layernorm_scale_shift_gate_select01_kernel(
+    output_ptr,
+    gate_out_ptr,
+    x_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale0_ptr,
+    shift0_ptr,
+    gate0_ptr,
+    scale1_ptr,
+    shift1_ptr,
+    gate1_ptr,
+    index_ptr,
+    inner_dim,
+    seq_len,
+    stride_x_row,
+    stride_out_row,
+    stride_go_row,
+    stride_w,
+    stride_b,
+    stride_s0_b,
+    stride_s0_c,
+    stride_sh0_b,
+    stride_sh0_c,
+    stride_g0_b,
+    stride_g0_c,
+    stride_s1_b,
+    stride_s1_c,
+    stride_sh1_b,
+    stride_sh1_c,
+    stride_g1_b,
+    stride_g1_c,
+    stride_i_b,
+    stride_i_l,
+    eps,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < inner_dim
+
+    x_row_ptr = x_ptr + row * stride_x_row
+    out_row_ptr = output_ptr + row * stride_out_row
+    gate_row_ptr = gate_out_ptr + row * stride_go_row
+
+    x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    mean = tl.sum(x, axis=0) / inner_dim
+    xbar = tl.where(mask, x - mean, 0.0)
+    var = tl.sum(xbar * xbar, axis=0) / inner_dim
+    rstd = tl.rsqrt(var + eps)
+    x_hat = (x - mean) * rstd
+
+    if HAS_WEIGHT:
+        w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+        x_hat = x_hat * w
+    if HAS_BIAS:
+        b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+        x_hat = x_hat + b
+
+    batch_idx = row // seq_len
+    seq_idx = row % seq_len
+    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+    scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
+    shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
+    gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
+
+    scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
+    shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
+    gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
+
+    # Branch on scalar idx instead of using tl.where on pointers.
+    # tl.where on pointers triggers an assertion in AMD Triton's
+    # CanonicalizePointers pass (ConvertArithSelectOp) on gfx950.
+    # This keeps it at 3 loads (not 6), avoids the pointer-level
+    # tl.where entirely, and since idx is uniform across all threads
+    # the branch has no divergence cost.
+    if idx:
+        scale = tl.load(scale1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate1_ptrs, mask=mask, other=0.0)
+    else:
+        scale = tl.load(scale0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate0_ptrs, mask=mask, other=0.0)
+    y = x_hat * (1.0 + scale) + shift
+
+    tl.store(out_row_ptr + cols, y, mask=mask)
+    tl.store(gate_row_ptr + cols, gate, mask=mask)
+
+
+@triton.jit
+def _fused_residual_layernorm_scale_shift_gate_select01_kernel(
+    output_ptr,
+    residual_out_ptr,
+    gate_out_ptr,
+    x_ptr,
+    residual_ptr,
+    residual_gate_ptr,
+    weight_ptr,
+    bias_ptr,
+    scale0_ptr,
+    shift0_ptr,
+    gate0_ptr,
+    scale1_ptr,
+    shift1_ptr,
+    gate1_ptr,
+    index_ptr,
+    inner_dim,
+    seq_len,
+    stride_x_row,
+    stride_res_row,
+    stride_rg_row,
+    stride_out_row,
+    stride_res_out_row,
+    stride_go_row,
+    stride_w,
+    stride_b,
+    stride_s0_b,
+    stride_s0_c,
+    stride_sh0_b,
+    stride_sh0_c,
+    stride_g0_b,
+    stride_g0_c,
+    stride_s1_b,
+    stride_s1_c,
+    stride_sh1_b,
+    stride_sh1_c,
+    stride_g1_b,
+    stride_g1_c,
+    stride_i_b,
+    stride_i_l,
+    eps,
+    HAS_WEIGHT: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_N)
+    mask = cols < inner_dim
+
+    x_row_ptr = x_ptr + row * stride_x_row
+    res_row_ptr = residual_ptr + row * stride_res_row
+    rg_row_ptr = residual_gate_ptr + row * stride_rg_row
+    out_row_ptr = output_ptr + row * stride_out_row
+    res_out_row_ptr = residual_out_ptr + row * stride_res_out_row
+    gate_row_ptr = gate_out_ptr + row * stride_go_row
+
+    x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual = tl.load(res_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual_gate = tl.load(rg_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    residual_out = residual + residual_gate * x
+    tl.store(res_out_row_ptr + cols, residual_out, mask=mask)
+
+    mean = tl.sum(residual_out, axis=0) / inner_dim
+    xbar = tl.where(mask, residual_out - mean, 0.0)
+    var = tl.sum(xbar * xbar, axis=0) / inner_dim
+    rstd = tl.rsqrt(var + eps)
+    x_hat = (residual_out - mean) * rstd
+
+    if HAS_WEIGHT:
+        w = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32)
+        x_hat = x_hat * w
+    if HAS_BIAS:
+        b = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32)
+        x_hat = x_hat + b
+
+    batch_idx = row // seq_len
+    seq_idx = row % seq_len
+    idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1)
+
+    scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c
+    shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c
+    gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c
+
+    scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c
+    shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c
+    gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c
+
+    # Branch on scalar idx instead of using tl.where on pointers.
+    # tl.where on pointers triggers an assertion in AMD Triton's
+    # CanonicalizePointers pass (ConvertArithSelectOp) on gfx950.
+    # This keeps it at 3 loads (not 6), avoids the pointer-level
+    # tl.where entirely, and since idx is uniform across all threads
+    # the branch has no divergence cost.
+    if idx:
+        scale = tl.load(scale1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift1_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate1_ptrs, mask=mask, other=0.0)
+    else:
+        scale = tl.load(scale0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        shift = tl.load(shift0_ptrs, mask=mask, other=0.0).to(tl.float32)
+        gate = tl.load(gate0_ptrs, mask=mask, other=0.0)
+    y = x_hat * (1.0 + scale) + shift
+
+    tl.store(out_row_ptr + cols, y, mask=mask)
+    tl.store(gate_row_ptr + cols, gate, mask=mask)
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_N": 64}, num_warps=2),
+        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.
+    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 (per-frame) and shift (per-token)
+    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 is per-token [B*L, C], indexed by pid_row directly
+    shift_ptrs = shift_ptr + pid_row * 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)
+
+
+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) or (x.is_xpu and scale.is_xpu)
+    assert x.is_contiguous()
+
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    if x.numel() == 0:
+        return output
+
+    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)
+
+        def grid(meta):
+            return (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 scale [B, F, 1, C] -> [B*F, C] (per-frame)
+        scale_reshaped = scale.squeeze(2).reshape(-1, C).contiguous()
+        if shift.dim() == 4 and current_platform.is_hip():
+            # ROCm has no fused CUTLASS scale-shift kernel, so this native path
+            # handles the causal Wan / LingBot output AdaLN, which passes a
+            # per-frame shift [B, F, 1, C]. Broadcast it across each frame's
+            # tokens to per-token [B, L, C] before flattening to [B*L, C],
+            # matching the per-token indexing in _fused_scale_shift_4d_kernel
+            # (the CUDA fused path accepts [B, F, 1, C] shift and broadcasts it
+            # per-frame).
+            shift_reshaped = (
+                shift.expand(B, num_frames, frame_seqlen, C)
+                .reshape(rows, C)
+                .contiguous()
+            )
+        else:
+            # shift is per-token [B, L, C] -> [B*L, C]
+            shift_reshaped = shift.reshape(rows, C).contiguous()
+
+        _fused_scale_shift_4d_kernel[grid](
+            output_2d,
+            x_2d,
+            scale_reshaped,
+            shift_reshaped,
+            scale_constant,
+            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 not (
+                scale_blc.any().to("cpu", non_blocking=True)
+                or shift_blc.any().to("cpu", non_blocking=True)
+            ):
+                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_layernorm_scale_shift_gate_select01_kernel(
+    x: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    scale0: torch.Tensor,
+    shift0: torch.Tensor,
+    gate0: torch.Tensor,
+    scale1: torch.Tensor,
+    shift1: torch.Tensor,
+    gate1: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+):
+    assert x.is_cuda
+    assert x.is_contiguous()
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    gate_out = torch.empty_like(x)
+
+    if (
+        scale0.dim() != 2
+        or shift0.dim() != 2
+        or gate0.dim() != 2
+        or scale1.dim() != 2
+        or shift1.dim() != 2
+        or gate1.dim() != 2
+    ):
+        raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
+    if index.dim() != 2:
+        raise ValueError("index must be 2D [B, L]")
+    if weight is not None and (weight.dim() != 1 or weight.shape[0] != C):
+        raise ValueError("weight must be 1D [C]")
+    if bias is not None and (bias.dim() != 1 or bias.shape[0] != C):
+        raise ValueError("bias must be 1D [C]")
+
+    x_2d = x.view(B * L, C)
+    output_2d = output.view(B * L, C)
+    gate_out_2d = gate_out.view(B * L, C)
+    weight = weight.contiguous() if weight is not None else x_2d
+    bias = bias.contiguous() if bias is not None else x_2d
+
+    MAX_FUSED_SIZE = 65536 // x_2d.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C))
+    if C > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    num_warps, num_stages = 4, 4
+
+    grid = (B * L,)
+    _fused_layernorm_scale_shift_gate_select01_kernel[grid](
+        output_2d,
+        gate_out_2d,
+        x_2d,
+        weight,
+        bias,
+        scale0.contiguous(),
+        shift0.contiguous(),
+        gate0.contiguous(),
+        scale1.contiguous(),
+        shift1.contiguous(),
+        gate1.contiguous(),
+        index.contiguous(),
+        C,
+        L,
+        x_2d.stride(0),
+        output_2d.stride(0),
+        gate_out_2d.stride(0),
+        weight.stride(0) if weight.dim() == 1 else 0,
+        bias.stride(0) if bias.dim() == 1 else 0,
+        scale0.stride(0),
+        scale0.stride(1),
+        shift0.stride(0),
+        shift0.stride(1),
+        gate0.stride(0),
+        gate0.stride(1),
+        scale1.stride(0),
+        scale1.stride(1),
+        shift1.stride(0),
+        shift1.stride(1),
+        gate1.stride(0),
+        gate1.stride(1),
+        index.stride(0),
+        index.stride(1),
+        eps,
+        HAS_WEIGHT=weight is not x_2d,
+        HAS_BIAS=bias is not x_2d,
+        BLOCK_N=BLOCK_N,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return output, gate_out
+
+
+def fuse_residual_layernorm_scale_shift_gate_select01_kernel(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    residual_gate: torch.Tensor,
+    weight: torch.Tensor | None,
+    bias: torch.Tensor | None,
+    scale0: torch.Tensor,
+    shift0: torch.Tensor,
+    gate0: torch.Tensor,
+    scale1: torch.Tensor,
+    shift1: torch.Tensor,
+    gate1: torch.Tensor,
+    index: torch.Tensor,
+    eps: float,
+):
+    assert x.is_cuda
+    assert x.is_contiguous()
+    assert residual.is_contiguous()
+    assert residual_gate.is_contiguous()
+    B, L, C = x.shape
+    output = torch.empty_like(x)
+    residual_out = torch.empty_like(x)
+    gate_out = torch.empty_like(x)
+
+    if residual.shape != x.shape:
+        raise ValueError("residual must have the same shape as x")
+    if residual_gate.shape != x.shape:
+        raise ValueError("residual_gate must have the same shape as x")
+    if (
+        scale0.dim() != 2
+        or shift0.dim() != 2
+        or gate0.dim() != 2
+        or scale1.dim() != 2
+        or shift1.dim() != 2
+        or gate1.dim() != 2
+    ):
+        raise ValueError("scale0/shift0/gate0/scale1/shift1/gate1 must be 2D [B, C]")
+    if index.dim() != 2:
+        raise ValueError("index must be 2D [B, L]")
+    if weight is not None and (weight.dim() != 1 or weight.shape[0] != C):
+        raise ValueError("weight must be 1D [C]")
+    if bias is not None and (bias.dim() != 1 or bias.shape[0] != C):
+        raise ValueError("bias must be 1D [C]")
+
+    x_2d = x.view(B * L, C)
+    residual_2d = residual.view(B * L, C)
+    residual_gate_2d = residual_gate.view(B * L, C)
+    output_2d = output.view(B * L, C)
+    residual_out_2d = residual_out.view(B * L, C)
+    gate_out_2d = gate_out.view(B * L, C)
+    weight = weight.contiguous() if weight is not None else x_2d
+    bias = bias.contiguous() if bias is not None else x_2d
+
+    MAX_FUSED_SIZE = 65536 // x_2d.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(C))
+    if C > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    num_warps, num_stages = 4, 4
+
+    grid = (B * L,)
+    _fused_residual_layernorm_scale_shift_gate_select01_kernel[grid](
+        output_2d,
+        residual_out_2d,
+        gate_out_2d,
+        x_2d,
+        residual_2d,
+        residual_gate_2d,
+        weight,
+        bias,
+        scale0.contiguous(),
+        shift0.contiguous(),
+        gate0.contiguous(),
+        scale1.contiguous(),
+        shift1.contiguous(),
+        gate1.contiguous(),
+        index.contiguous(),
+        C,
+        L,
+        x_2d.stride(0),
+        residual_2d.stride(0),
+        residual_gate_2d.stride(0),
+        output_2d.stride(0),
+        residual_out_2d.stride(0),
+        gate_out_2d.stride(0),
+        weight.stride(0) if weight.dim() == 1 else 0,
+        bias.stride(0) if bias.dim() == 1 else 0,
+        scale0.stride(0),
+        scale0.stride(1),
+        shift0.stride(0),
+        shift0.stride(1),
+        gate0.stride(0),
+        gate0.stride(1),
+        scale1.stride(0),
+        scale1.stride(1),
+        shift1.stride(0),
+        shift1.stride(1),
+        gate1.stride(0),
+        gate1.stride(1),
+        index.stride(0),
+        index.stride(1),
+        eps,
+        HAS_WEIGHT=weight is not x_2d,
+        HAS_BIAS=bias is not x_2d,
+        BLOCK_N=BLOCK_N,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return output, residual_out, gate_out
+
+
+if current_platform.is_npu():
+    from .npu_fallback import fuse_scale_shift_native
+
+    fuse_scale_shift_kernel = fuse_scale_shift_native
+
+if current_platform.is_mps():
+    from .mps_fallback import fuse_scale_shift_kernel_native
+
+    fuse_scale_shift_kernel = fuse_scale_shift_kernel_native
+
+if current_platform.is_musa():
+    from .torch_fallback import fuse_scale_shift_kernel_native
+
+    fuse_scale_shift_kernel = fuse_scale_shift_kernel_native
+
+if current_platform.is_cpu():
+    from .torch_fallback import (
+        fuse_scale_shift_kernel_native,
+    )
+
+    fuse_scale_shift_kernel = fuse_scale_shift_kernel_native

python/sglang/jit_kernel/diffusion/triton/torch_fallback.py

@@ -0,0 +1,146 @@
+"""Pytorch native based fallbacks for Triton diffusion kernels.
+
+Triton is not available on some platforms, so these pure-PyTorch
+implementations replace the Triton kernels
+
+"""
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+def fuse_scale_shift_kernel_native(
+    x: torch.Tensor,
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    scale_constant: float = 1.0,
+    block_l: int = 128,
+    block_c: int = 128,
+):
+    """Native fallback for fuse_scale_shift_kernel with scale_constant support."""
+    B, L, C = x.shape
+
+    def _expand(t: torch.Tensor) -> torch.Tensor:
+        if t.dim() == 4:
+            # [B, F, 1, C] -> [B, L, C]
+            num_frames = t.shape[1]
+            frame_seqlen = L // num_frames
+            return (
+                t.squeeze(2)
+                .unsqueeze(2)
+                .expand(-1, -1, frame_seqlen, -1)
+                .reshape(B, L, C)
+            )
+        elif t.dim() == 2:
+            # [B, C] -> [B, 1, C]
+            return t.unsqueeze(1)
+        return t
+
+    scale = _expand(scale)
+    shift = _expand(shift)
+
+    return x * (scale_constant + scale) + shift
+
+
+def apply_rotary_embedding_native(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
+) -> torch.Tensor:
+    """Native fallback for rotary embedding (shared with NPU implementation)."""
+    if interleaved and cos.shape[-1] == x.shape[-1]:
+        cos = cos[..., ::2]
+        sin = sin[..., ::2]
+    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)
+
+
+def norm_infer_native(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    eps: float,
+    is_rms_norm: bool = False,
+    out: Optional[Tensor] = None,
+) -> Tensor:
+    """Native fallback for norm_infer (layer norm / rms norm inference)."""
+    orig_dtype = x.dtype
+    x = x.contiguous().float()
+    if is_rms_norm:
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x_hat = x * torch.rsqrt(variance + eps)
+    else:
+        mean = x.mean(dim=-1, keepdim=True)
+        variance = (x - mean).pow(2).mean(dim=-1, keepdim=True)
+        x_hat = (x - mean) * torch.rsqrt(variance + eps)
+    if weight is not None:
+        x_hat = x_hat * weight.float()
+    if bias is not None:
+        x_hat = x_hat + bias.float()
+    result = x_hat.to(orig_dtype)
+    if out is not None:
+        out.copy_(result)
+        return out
+    return result
+
+
+def triton_one_pass_rms_norm_native(
+    x: torch.Tensor, w: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    """Native fallback for triton_one_pass_rms_norm."""
+    shape = x.shape
+    orig_dtype = x.dtype
+    x = x.contiguous().float()
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x_hat = x * torch.rsqrt(variance + eps)
+    return (x_hat * w.float()).to(orig_dtype).view(shape)
+
+
+def rms_norm_fn_native(
+    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,
+):
+    """Native fallback for rms_norm_fn (inference only, no dropout/x1 support)."""
+    x_shape_og = x.shape
+    orig_dtype = x.dtype
+    x = x.reshape(-1, x.shape[-1]).float()
+    if residual is not None:
+        residual = residual.reshape(-1, residual.shape[-1]).float()
+        x = x + residual
+        residual_out_val = x.to(torch.float32 if residual_in_fp32 else orig_dtype)
+    else:
+        residual_out_val = None
+    variance = x.pow(2).mean(dim=-1, keepdim=True)
+    x_hat = x * torch.rsqrt(variance + eps)
+    if weight is not None:
+        w = weight.float()
+        if zero_centered_weight:
+            w = w + 1.0
+        x_hat = x_hat * w
+    if bias is not None:
+        x_hat = x_hat + bias.float()
+    final_dtype = out_dtype if out_dtype is not None else orig_dtype
+    y = x_hat.to(final_dtype).reshape(x_shape_og)
+    if residual is not None and residual_out_val is not None:
+        return y, residual_out_val.reshape(x_shape_og)
+    return y

python/sglang/jit_kernel/diffusion/triton/varlen_pack_pad.py

@@ -0,0 +1,191 @@
+"""Fused Triton pack/scatter kernels for the varlen mask path.
+
+Used by ``USPAttention.forward`` masked branch to gather Q/K/V at valid
+positions and scatter the FA output back to the dense ``[B, S, H, D]`` layout.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+# ---------------------------------------------------------------------------
+# Pack (unpad) — gather Q/K/V at indices into packed [total_valid, H, D]
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _fused_pack_qkv_kernel(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    Q_unpad_ptr,
+    K_unpad_ptr,
+    V_unpad_ptr,
+    indices_ptr,
+    HD,  # H * D, flattened feature dim
+    src_row_stride,  # stride between rows in Q/K/V (B*S row -> next row)
+    dst_row_stride,  # stride in Q_unpad/K_unpad/V_unpad
+    BLOCK_HD: tl.constexpr,
+):
+    """One program per packed row; copies Q[src], K[src], V[src] to dst row."""
+    out_row = tl.program_id(0)
+    src_row = tl.load(indices_ptr + out_row).to(tl.int64)
+
+    cols = tl.arange(0, BLOCK_HD)
+    col_mask = cols < HD
+
+    src_offset = src_row * src_row_stride + cols
+    dst_offset = out_row * dst_row_stride + cols
+
+    q_val = tl.load(Q_ptr + src_offset, mask=col_mask)
+    k_val = tl.load(K_ptr + src_offset, mask=col_mask)
+    v_val = tl.load(V_ptr + src_offset, mask=col_mask)
+
+    tl.store(Q_unpad_ptr + dst_offset, q_val, mask=col_mask)
+    tl.store(K_unpad_ptr + dst_offset, k_val, mask=col_mask)
+    tl.store(V_unpad_ptr + dst_offset, v_val, mask=col_mask)
+
+
+def fused_pack_qkv(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    indices: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Pack ``[B, S, H, D]`` Q/K/V at ``indices`` into ``[total_valid, H, D]``.
+
+    ``indices`` is the int64 flat ``B*S`` position for each kept token.
+    Non-contiguous inputs are made contiguous internally.
+    """
+    assert q.shape == k.shape == v.shape, "Q/K/V must share shape"
+    assert q.dtype == k.dtype == v.dtype, "Q/K/V must share dtype"
+    assert q.dim() == 4, "Q/K/V must be [B, S, H, D]"
+    assert indices.dtype in (torch.int32, torch.int64)
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    bs, seq, num_heads, head_dim = q.shape
+    hd = num_heads * head_dim
+    n_valid = indices.shape[0]
+
+    if n_valid == 0:
+        return (
+            q.new_empty(0, num_heads, head_dim),
+            k.new_empty(0, num_heads, head_dim),
+            v.new_empty(0, num_heads, head_dim),
+        )
+
+    block_hd = triton.next_power_of_2(hd)
+    q_flat = q.view(bs * seq, hd)
+    k_flat = k.view(bs * seq, hd)
+    v_flat = v.view(bs * seq, hd)
+    q_unpad = torch.empty(n_valid, hd, dtype=q.dtype, device=q.device)
+    k_unpad = torch.empty(n_valid, hd, dtype=k.dtype, device=k.device)
+    v_unpad = torch.empty(n_valid, hd, dtype=v.dtype, device=v.device)
+
+    with torch.get_device_module().device(q.device):
+        _fused_pack_qkv_kernel[(n_valid,)](
+            q_flat,
+            k_flat,
+            v_flat,
+            q_unpad,
+            k_unpad,
+            v_unpad,
+            indices,
+            hd,
+            q_flat.stride(0),
+            q_unpad.stride(0),
+            BLOCK_HD=block_hd,
+        )
+
+    return (
+        q_unpad.view(n_valid, num_heads, head_dim),
+        k_unpad.view(n_valid, num_heads, head_dim),
+        v_unpad.view(n_valid, num_heads, head_dim),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Scatter (pad) — write packed output to [B, S, H, D] with zeros at invalid
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _fused_scatter_to_padded_kernel(
+    Out_unpad_ptr,
+    Out_padded_ptr,
+    inv_indices_ptr,  # [B*S]: pack idx for valid row, -1 for invalid
+    HD,
+    src_row_stride,
+    dst_row_stride,
+    BLOCK_HD: tl.constexpr,
+):
+    """One program per padded row; writes from pack or zeros."""
+    pad_row = tl.program_id(0)
+    inv_idx = tl.load(inv_indices_ptr + pad_row).to(tl.int64)
+
+    cols = tl.arange(0, BLOCK_HD)
+    col_mask = cols < HD
+    valid = inv_idx >= 0
+
+    safe_idx = tl.where(valid, inv_idx, 0)
+    src_offset = safe_idx * src_row_stride + cols
+    dst_offset = pad_row * dst_row_stride + cols
+
+    val = tl.load(Out_unpad_ptr + src_offset, mask=col_mask & valid, other=0.0)
+    tl.store(Out_padded_ptr + dst_offset, val, mask=col_mask)
+
+
+def fused_scatter_to_padded(
+    out_unpad: torch.Tensor,
+    inv_indices: torch.Tensor,
+    batch_size: int,
+    seqlen: int,
+) -> torch.Tensor:
+    """Scatter packed varlen output back to ``[B, S, H, D]`` with zero padding.
+
+    ``inv_indices`` is ``[B*S]`` giving the pack row index for each padded
+    position (``-1`` for padding). Non-contiguous ``out_unpad`` is made contiguous.
+    """
+    assert out_unpad.dim() == 3, "out_unpad must be [total_valid, H, D]"
+    assert inv_indices.shape == (batch_size * seqlen,)
+    assert inv_indices.dtype in (torch.int32, torch.int64)
+    out_unpad = out_unpad.contiguous()
+    _, num_heads, head_dim = out_unpad.shape
+    hd = num_heads * head_dim
+    block_hd = triton.next_power_of_2(hd)
+
+    out_padded = torch.empty(
+        batch_size * seqlen, hd, dtype=out_unpad.dtype, device=out_unpad.device
+    )
+    out_unpad_flat = out_unpad.view(-1, hd)
+
+    with torch.get_device_module().device(out_unpad.device):
+        _fused_scatter_to_padded_kernel[(batch_size * seqlen,)](
+            out_unpad_flat,
+            out_padded,
+            inv_indices,
+            hd,
+            out_unpad_flat.stride(0),
+            out_padded.stride(0),
+            BLOCK_HD=block_hd,
+        )
+
+    return out_padded.view(batch_size, seqlen, num_heads, head_dim)
+
+
+# ---------------------------------------------------------------------------
+# Inverse-index builder (called once per request alongside indices)
+# ---------------------------------------------------------------------------
+
+
+def build_inv_indices(indices: torch.Tensor, total_rows: int) -> torch.Tensor:
+    """For each padded row in ``[B*S]``, return its pack index or ``-1``."""
+    n_valid = indices.shape[0]
+    inv = torch.full((total_rows,), -1, dtype=torch.int32, device=indices.device)
+    inv[indices.long()] = torch.arange(
+        n_valid, dtype=torch.int32, device=indices.device
+    )
+    return inv

python/sglang/jit_kernel/diffusion/triton/zimage_native_norm.py

@@ -0,0 +1,182 @@
+import torch
+import triton  # type: ignore
+import triton.language as tl  # type: ignore
+
+
+@triton.jit
+def _tanh(x):
+    return 2.0 / (1.0 + tl.exp(-2.0 * x)) - 1.0
+
+
+@triton.jit
+def _rmsnorm_scale_kernel(
+    y_ptr,
+    x_ptr,
+    weight_ptr,
+    scale_ptr,
+    x_row_stride,
+    scale_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    scale = tl.load(
+        scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0
+    )
+    y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+@triton.jit
+def _rmsnorm_tanh_residual_kernel(
+    y_ptr,
+    x_ptr,
+    gate_ptr,
+    residual_ptr,
+    weight_ptr,
+    x_row_stride,
+    gate_row_stride,
+    residual_row_stride,
+    seq_len,
+    dim: tl.constexpr,
+    eps: tl.constexpr,
+    block_dim: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, block_dim)
+    mask = offsets < dim
+
+    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
+    square = (x * x).to(tl.bfloat16)
+    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
+    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
+
+    batch = row // seq_len
+    gate = tl.load(gate_ptr + batch * gate_row_stride + offsets, mask=mask, other=0.0)
+    residual = tl.load(
+        residual_ptr + row * residual_row_stride + offsets, mask=mask, other=0.0
+    )
+    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
+    norm = ((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16)
+    gated = (_tanh(gate.to(tl.float32)).to(tl.bfloat16) * norm).to(tl.bfloat16)
+    y = (residual + gated).to(tl.bfloat16)
+    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
+
+
+def _flat_row_stride(x: torch.Tensor) -> int | None:
+    if x.dim() < 2 or x.stride(-1) != 1:
+        return None
+    row_stride = x.stride(-2)
+    expected_stride = row_stride * x.shape[-2]
+    for dim in range(x.dim() - 3, -1, -1):
+        if x.stride(dim) != expected_stride:
+            return None
+        expected_stride *= x.shape[dim]
+    return row_stride
+
+
+def _can_use(x: torch.Tensor, weight: torch.Tensor, other: torch.Tensor) -> bool:
+    return (
+        x.is_cuda
+        and weight.is_cuda
+        and other.is_cuda
+        and x.dtype == torch.bfloat16
+        and weight.dtype == torch.bfloat16
+        and other.dtype == torch.bfloat16
+        and weight.is_contiguous()
+        and x.shape[-1] <= 8192
+        and _flat_row_stride(x) is not None
+        and _flat_row_stride(other) is not None
+    )
+
+
+def zimage_rmsnorm_scale(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    scale: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    if not _can_use(x, weight, scale):
+        return None
+    shape = x.shape
+    dim = shape[-1]
+    x_rows = x.numel() // dim
+    scale_rows = scale.numel() // dim
+    if x_rows % scale_rows != 0:
+        return None
+    seq_len = x_rows // scale_rows
+    x_row_stride = _flat_row_stride(x)
+    scale_row_stride = _flat_row_stride(scale)
+    if x_row_stride is None or scale_row_stride is None:
+        return None
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_scale_kernel[(x_rows,)](
+            y.reshape(-1, dim),
+            x,
+            weight,
+            scale,
+            x_row_stride,
+            scale_row_stride,
+            seq_len,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return y
+
+
+def zimage_rmsnorm_tanh_residual(
+    x: torch.Tensor,
+    gate: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor | None:
+    if not (_can_use(x, weight, gate) and residual.is_cuda):
+        return None
+    if residual.dtype != x.dtype or _flat_row_stride(residual) is None:
+        return None
+    shape = x.shape
+    dim = shape[-1]
+    x_rows = x.numel() // dim
+    gate_rows = gate.numel() // dim
+    if x_rows % gate_rows != 0:
+        return None
+    seq_len = x_rows // gate_rows
+    x_row_stride = _flat_row_stride(x)
+    gate_row_stride = _flat_row_stride(gate)
+    residual_row_stride = _flat_row_stride(residual)
+    if x_row_stride is None or gate_row_stride is None or residual_row_stride is None:
+        return None
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    with torch.get_device_module().device(x.device):
+        _rmsnorm_tanh_residual_kernel[(x_rows,)](
+            y.reshape(-1, dim),
+            x,
+            gate,
+            residual,
+            weight,
+            x_row_stride,
+            gate_row_stride,
+            residual_row_stride,
+            seq_len,
+            dim,
+            eps,
+            block_dim=triton.next_power_of_2(dim),
+            num_warps=8,
+        )
+    return y