chore: import upstream snapshot with attribution

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)