[diffusion] kernel: gated residual layernorm scale shift and layernorm scale shift kernel fusion for Qwen-Image, WAN and HunyuanVideo (#14717)

Co-authored-by: AichenF <aichenf@nvidia.com>
Co-authored-by: jianyingzhu <joeyzhu@nvidia.com>
Co-authored-by: root <root@a4u8g-0120.ipp2a2.colossus.nvidia.com>
Co-authored-by: Yihan Chen <yingluosanqian@example.com>
Co-authored-by: 陈一涵 <yingluosanqian@gmail.com>
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py

@@ -0,0 +1,134 @@
+# Benchmarks SGLang fused layernorm/rmsnorm scale shift kernels
+# 1. fused_norm_scale_shift
+# 2. fused_scale_residual_norm_scale_shift
+import itertools
+from typing import Tuple
+
+import torch
+import triton
+import triton.testing
+
+from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    LayerNormScaleShift,
+    RMSNormScaleShift,
+    ScaleResidualLayerNormScaleShift,
+    ScaleResidualRMSNormScaleShift,
+)
+
+if is_in_ci():
+    B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024]
+else:
+    B_RANGE, S_RANGE, D_RANGE = [1], [128, 1024, 4096], [1024, 3072, 4096]
+
+NORM_TYPE_RANGE = ["layer", "rms"]
+AFFINE_RANGE = [True, False]
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+EPS = 1e-5
+LINE_VALS = ["native", "cuda"]
+LINE_NAMES = ["SGLang Native", "SGLang Fused"]
+STYLES = [("red", "-"), ("blue", "--")]
+config = list(
+    itertools.product(B_RANGE, S_RANGE, D_RANGE, NORM_TYPE_RANGE, AFFINE_RANGE)
+)
+
+
+def preprocess_layer(layer, affine: bool, D: int, DTYPE: torch.dtype):
+    if affine:
+        weight = torch.randn(D, dtype=DTYPE, device=DEVICE)
+        bias = torch.randn(D, dtype=DTYPE, device=DEVICE)
+        with torch.no_grad():
+            layer.norm.weight.copy_(weight)
+            if hasattr(layer.norm, "bias"):
+                layer.norm.bias.copy_(bias)
+    layer.requires_grad_(False)
+    return layer.to(DEVICE)
+
+
+# ============================================================================
+# Benchmark 1: fused_norm_scale_shift
+# ============================================================================
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["B", "S", "D", "norm_type", "affine"],
+        x_vals=config,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=STYLES,
+        ylabel="us",
+        plot_name="fused_norm_scale_shift",
+        args={},
+    )
+)
+def bench_fused_norm_scale_shift(
+    B: int, S: int, D: int, norm_type, affine: bool, provider: str
+) -> Tuple[float, float, float]:
+    x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    if norm_type == "layer":
+        layer = LayerNormScaleShift(D, EPS, affine, dtype=DTYPE)
+    else:
+        layer = RMSNormScaleShift(D, EPS, affine, dtype=DTYPE)
+    layer = preprocess_layer(layer, affine, D, DTYPE)
+    if provider == "native":
+        fn = lambda: layer.forward_native(x, shift, scale)
+    else:
+        fn = lambda: layer.forward_cuda(x, shift, scale)
+
+    quantiles = [0.5, 0.2, 0.8]
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms  # convert to us
+
+
+# ============================================================================
+# Benchmark 2: fused_scale_residual_norm_scale_shift
+# ============================================================================
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["B", "S", "D", "norm_type", "affine"],
+        x_vals=config,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=STYLES,
+        ylabel="us",
+        plot_name="fused_scale_residual_norm_scale_shift",
+        args={},
+    )
+)
+def bench_fused_scale_residual_norm_scale_shift(
+    B: int, S: int, D: int, norm_type, affine: bool, provider: str
+) -> Tuple[float, float, float]:
+    residual = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    scale = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    shift = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    gate = torch.randn(B, 1, D, dtype=DTYPE, device=DEVICE)
+    if norm_type == "layer":
+        layer = ScaleResidualLayerNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE)
+    else:
+        layer = ScaleResidualRMSNormScaleShift(D, EPS, affine, dtype=DTYPE).to(DEVICE)
+    layer = preprocess_layer(layer, affine, D, DTYPE)
+    if provider == "native":
+        fn = lambda: layer.forward_native(residual, x, gate, shift, scale)
+    else:
+        fn = lambda: layer.forward_cuda(residual, x, gate, shift, scale)
+
+    quantiles = [0.5, 0.2, 0.8]
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms  # convert to us
+
+
+if __name__ == "__main__":
+    print(f"\n{'='*80}")
+    print("Benchmark: fused_norm_scale_shift")
+    print(f"{'='*80}\n")
+    bench_fused_norm_scale_shift.run(print_data=True)
+
+    print(f"\n{'='*80}")
+    print("Benchmark: fused_scale_residual_norm_scale_shift")
+    print(f"{'='*80}\n")
+    bench_fused_scale_residual_norm_scale_shift.run(print_data=True)

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)
+    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[0] = val
+    cute.arch.sync_threads()
+    val = acc[0]
+    return val

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

@@ -0,0 +1,419 @@
+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 TORCH_TO_CUTE_DTYPE, WARP_SIZE
+
+_COMPILE_CACHE = {}
+
+
+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)
+
+
+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(src, 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
+        elif cutlass.const_expr(isinstance(tGrG, cutlass.Int32)):
+            value = tGrG * 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], B: int, S: int, 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 and (t.shape[0] != B or t.shape[2] != 1 or t.shape[3] != D):
+        F = t.shape[1]
+        if S % F != 0:
+            raise ValueError(f"Validate failed: S({S}) must be divisible by F({F}).")
+        failed = True
+    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._dynamo.disable  # Disable Dynamo tracing
+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,
+    stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
+) -> Tuple[torch.Tensor, 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`).
+    """
+    # Tensor Validation
+    BSD = x.shape
+    validate_x(x, *BSD)
+    validate_weight_bias(weight, *BSD)
+    validate_weight_bias(bias, *BSD)
+    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. Unless explicitly handled
+        # (e.g., for gate), 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]
+        cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
+        # 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(*cute_tensor_args, eps, stream)
+        return y
+    else:
+        raise ValueError(f'norm_type must be one of "layer" and "rms"')
+
+
+@torch._dynamo.disable  # Disable Dynamo tracing
+def fused_scale_residual_norm_scale_shift(
+    residual: torch.Tensor,
+    x: torch.Tensor,
+    gate: Union[Optional[torch.Tensor], int],
+    weight: Optional[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    scale: torch.Tensor,
+    shift: torch.Tensor,
+    norm_type: str,
+    eps: float = 1e-5,
+    stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
+) -> 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, [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`).
+    """
+    # Tensor Validation
+    BSD = x.shape
+    validate_x(x, *BSD)
+    validate_x(residual, *BSD)
+    validate_gate(gate, *BSD)
+    validate_weight_bias(weight, *BSD)
+    validate_weight_bias(bias, *BSD)
+    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
+        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. Unless explicitly handled
+        # (e.g., for gate), 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]
+        cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
+        # 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(*cute_tensor_args, eps, stream)
+        return y, resi_out
+    else:
+        raise ValueError(f'norm_type must be one of "layer" and "rms"')

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

@@ -0,0 +1,10 @@
+import cutlass
+import torch
+
+WARP_SIZE = 32
+
+TORCH_TO_CUTE_DTYPE = {
+    torch.float16: cutlass.Float16,
+    torch.bfloat16: cutlass.BFloat16,
+    torch.float32: cutlass.Float32,
+}

python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py

@@ -0,0 +1,238 @@
+from typing import Optional, Tuple
+
+import pytest
+import torch
+from einops import rearrange
+from torch import Tensor
+
+from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import (
+    fused_norm_scale_shift,
+    fused_scale_residual_norm_scale_shift,
+)
+
+DEVICE = "cuda"
+SHAPE_MAP = {
+    "1": lambda B, S, F, D: (1,),
+    "D": lambda B, S, F, D: (D,),
+    "1D": lambda B, S, F, D: (1, D),
+    "BD": lambda B, S, F, D: (B, D),
+    "11D": lambda B, S, F, D: (1, 1, D),
+    "B1D": lambda B, S, F, D: (B, 1, D),
+    "1SD": lambda B, S, F, D: (1, S, D),
+    "BSD": lambda B, S, F, D: (B, S, D),
+    "BF1D": lambda B, S, F, D: (B, F, 1, D),
+}
+SHAPES = [
+    # (B, S, F, D)
+    (1, 1024, 8, 3072),
+    (4, 512, 16, 3072),
+    (1, 115200, 1, 3072),  # Hunyuan
+    (1, 32760, 1, 1536),  # Wan
+    (1, 6, 1, 3072),  # Qwen
+]
+DTYPES = [torch.float16, torch.bfloat16, torch.float32]
+NORM_TYPES = ["layer", "rms"]
+AFFINE_MODES = ["D", "NAT"]
+INDEX_MODES = ["BSD", "1", "1SD", "BD", "B1D", "D", "1D", "11D", "BF1D"]
+
+
+def _tol(dtype: torch.dtype):
+    return 1e-5 if dtype == torch.float32 else 5e-2
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+def _apply_scale_shift(y: Tensor, scale: Tensor, shift: Tensor) -> Tensor:
+    if scale.ndim == 4:
+        num_frame = scale.shape[1]
+        return rearrange(
+            rearrange(y, "b (f l) d -> b f l d", f=num_frame) * (1 + scale) + shift,
+            "b f l d -> b (f l) d",
+        )
+    else:
+        scale = rearrange(scale, "b d -> b 1 d") if scale.ndim == 2 else scale
+        shift = rearrange(shift, "b d -> b 1 d") if shift.ndim == 2 else shift
+        return y * (1 + scale) + shift
+
+
+def fused_norm_scale_shift_ref(
+    x: Tensor,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    scale: Tensor,
+    shift: Tensor,
+    norm_type: str,
+    eps: float,
+) -> Tensor:
+    original_dtype = x.dtype
+    x, weight, bias, scale, shift = (
+        v.float() if v is not None else v for v in [x, weight, bias, scale, shift]
+    )
+    if norm_type == "layer":
+        norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
+    else:
+        norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
+    return _apply_scale_shift(norm, scale, shift).to(original_dtype)
+
+
+def fused_scale_residual_norm_scale_shift_ref(
+    residual: Tensor,
+    x: Tensor,
+    gate: Optional[Tensor] | int,
+    weight: Optional[Tensor],
+    bias: Optional[Tensor],
+    scale: Tensor,
+    shift: Tensor,
+    norm_type: str,
+    eps: float,
+):
+    original_dtype = x.dtype
+    residual, x, gate, weight, bias, scale, shift = (
+        v.float() if isinstance(v, Tensor) else v
+        for v in [residual, x, gate, weight, bias, scale, shift]
+    )
+    if isinstance(gate, int):
+        x = residual + gate * x
+    else:
+        if gate.ndim == 4:
+            num_frame = gate.shape[1]
+            x_fld = rearrange(x, "b (f l) d -> b f l d", f=num_frame)
+            x = residual + rearrange(x_fld * gate, "b f l d -> b (f l) d")
+        else:
+            gate = rearrange(gate, "b d -> b 1 d") if gate.ndim == 2 else gate
+            x = residual + gate * x
+    if norm_type == "layer":
+        norm = torch.layer_norm(x, x.shape[-1:], eps=eps, weight=weight, bias=bias)
+    else:
+        norm = torch.rms_norm(x, x.shape[-1:], eps=eps, weight=weight)
+    y_ref = _apply_scale_shift(norm, scale, shift)
+    return y_ref.to(original_dtype), x.to(original_dtype)
+
+
+def _make_tensor(index_mode: str, shape: Tuple, dtype: torch.dtype):
+    if index_mode == "int1":
+        return 1
+    if index_mode == "NAT":
+        return None
+    return torch.randn(*SHAPE_MAP[index_mode](*shape), device=DEVICE, dtype=dtype)
+
+
+@torch.no_grad()
+def run_norm_scale_shift(
+    shape=SHAPES[0],
+    dtype=DTYPES[0],
+    affine_dtype=DTYPES[0],
+    scale_dtype=DTYPES[0],
+    shift_dtype=DTYPES[0],
+    norm_type=NORM_TYPES[0],
+    affine_mode=AFFINE_MODES[0],
+    scale_mode="BSD",
+    shift_mode="BSD",
+    eps=1e-5,
+):
+    x = _make_tensor("BSD", shape, dtype)
+    weight = _make_tensor(affine_mode, shape, affine_dtype)
+    bias = _make_tensor(affine_mode, shape, affine_dtype)
+    scale = _make_tensor(scale_mode, shape, scale_dtype)
+    shift = _make_tensor(shift_mode, shape, shift_dtype)
+    y_dev = fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps)
+    y_ref = fused_norm_scale_shift_ref(x, weight, bias, scale, shift, norm_type, eps)
+    torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
+
+
+@torch.no_grad()
+def run_scale_resi_norm_scale_shift(
+    shape=SHAPES[0],
+    dtype=DTYPES[0],
+    affine_dtype=DTYPES[0],
+    scale_dtype=DTYPES[0],
+    shift_dtype=DTYPES[0],
+    norm_type=NORM_TYPES[0],
+    affine_mode=AFFINE_MODES[0],
+    gate_mode="B1D",
+    scale_mode="BSD",
+    shift_mode="BSD",
+    eps=1e-5,
+):
+    residual = _make_tensor("BSD", shape, dtype)
+    x = _make_tensor("BSD", shape, dtype)
+    gate = _make_tensor(gate_mode, shape, dtype)
+    weight = _make_tensor(affine_mode, shape, affine_dtype)
+    bias = _make_tensor(affine_mode, shape, affine_dtype)
+    scale = _make_tensor(scale_mode, shape, scale_dtype)
+    shift = _make_tensor(shift_mode, shape, shift_dtype)
+    y_dev, res_dev = fused_scale_residual_norm_scale_shift(
+        residual, x, gate, weight, bias, scale, shift, norm_type, eps
+    )
+    y_ref, res_ref = fused_scale_residual_norm_scale_shift_ref(
+        residual, x, gate, weight, bias, scale, shift, norm_type, eps
+    )
+    torch.testing.assert_close(y_dev, y_ref, atol=_tol(dtype), rtol=_tol(dtype))
+    torch.testing.assert_close(res_dev, res_ref, atol=_tol(dtype), rtol=_tol(dtype))
+
+
+@pytest.mark.parametrize("norm_type", NORM_TYPES)
+class TestFusedNormScaleShift:
+    @pytest.mark.parametrize("shape", SHAPES)
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_shape_dtype(self, shape, dtype, norm_type):
+        run_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
+
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_dtype_0(self, dtype, norm_type):
+        run_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
+
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_dtype_1(self, dtype, norm_type):
+        run_norm_scale_shift(scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type)
+
+    @pytest.mark.parametrize("affine_mode", AFFINE_MODES)
+    def test_normtype_affine(self, affine_mode, norm_type):
+        run_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
+
+    @pytest.mark.parametrize("index_mode", INDEX_MODES)
+    def test_index_mode(self, index_mode, norm_type):
+        run_norm_scale_shift(
+            scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
+        )
+
+
+@pytest.mark.parametrize("norm_type", NORM_TYPES)
+class TestFusedScaleResidualNormScaleShift:
+    @pytest.mark.parametrize("shape", SHAPES)
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_shape_dtype(self, shape, dtype, norm_type):
+        run_scale_resi_norm_scale_shift(shape=shape, dtype=dtype, norm_type=norm_type)
+
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_dtype_0(self, dtype, norm_type):
+        run_scale_resi_norm_scale_shift(affine_dtype=dtype, norm_type=norm_type)
+
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_dtype_1(self, dtype, norm_type):
+        run_scale_resi_norm_scale_shift(
+            scale_dtype=dtype, shift_dtype=dtype, norm_type=norm_type
+        )
+
+    @pytest.mark.parametrize("affine_mode", AFFINE_MODES)
+    def test_normtype_affine(self, affine_mode, norm_type):
+        run_scale_resi_norm_scale_shift(affine_mode=affine_mode, norm_type=norm_type)
+
+    @pytest.mark.parametrize("index_mode", INDEX_MODES)
+    def test_scale_shift_index_mode(self, index_mode, norm_type):
+        run_scale_resi_norm_scale_shift(
+            scale_mode=index_mode, shift_mode=index_mode, norm_type=norm_type
+        )
+
+    @pytest.mark.parametrize("index_mode", INDEX_MODES + ["int1"])
+    def test_gate_index_mode(self, index_mode, norm_type):
+        run_scale_resi_norm_scale_shift(gate_mode=index_mode, norm_type=norm_type)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -51,7 +51,7 @@ class RMSNorm(CustomOp):
         var_hidden_size: Optional[int] = None,
     ) -> None:
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
         self.variance_epsilon = eps
         self.hidden_size = hidden_size
         self.variance_size_override = (
@@ -71,6 +71,7 @@ class RMSNorm(CustomOp):
         residual: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         shape = x.shape
+        device = x.device
         x = x.reshape(-1, shape[-1])
         if residual is not None:
             residual_shape = residual.shape
@@ -249,56 +250,55 @@ class LayerNorm(CustomOp):
 class FP32LayerNorm(nn.LayerNorm):
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         origin_dtype = inputs.dtype
+        device = inputs.device
         return F.layer_norm(
             inputs.float(),
             self.normalized_shape,
-            self.weight.float() if self.weight is not None else None,
+            self.weight.float().to(device=device) if self.weight is not None else None,
-            self.bias.float() if self.bias is not None else None,
+            self.bias.float().to(device=device) if self.bias is not None else None,
             self.eps,
         ).to(origin_dtype)
 
 
-class ScaleResidualLayerNormScaleShift(nn.Module):
+################################################################################
-    """
+# Fused norm kernel
-    Fused operation that combines:
+################################################################################
-    1. Gated residual connection
+def _ensure_contiguous(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
-    2. LayerNorm
+    return tensor.contiguous() if tensor is not None else None
-    3. Scale and shift operations
 
-    This reduces memory bandwidth by combining memory-bound operations.
+
+class _ScaleResidualNormScaleShift(CustomOp):
     """
+    Fused kernel that combines:
+    1. residual_out = residual + gate * x
+    2. normed = layernorm(residual_out) or rmsnorm(residual_out)
+    3. out = normed * (1 + scale) + shift
+    compute_dtype is always fp32 for higher precision.
+    """
+
+    norm_type: str
 
     def __init__(
         self,
         hidden_size: int,
-        norm_type: str = "rms",
         eps: float = 1e-6,
         elementwise_affine: bool = False,
         dtype: torch.dtype = torch.float32,
-        compute_dtype: torch.dtype | None = None,
         prefix: str = "",
     ):
         super().__init__()
-        if norm_type == "rms":
+        self.eps = eps
-            self.norm = RMSNorm(
+        self.dtype = dtype
-                hidden_size, has_weight=elementwise_affine, eps=eps, dtype=dtype
+        if self.norm_type == "rms":
-            )
+            self.norm = RMSNorm(hidden_size, eps=eps, dtype=dtype)
-        elif norm_type == "layer":
+        elif self.norm_type == "layer":
-            if compute_dtype == torch.float32:
             self.norm = FP32LayerNorm(
-                    hidden_size, elementwise_affine=elementwise_affine, eps=eps
+                hidden_size, elementwise_affine=elementwise_affine, eps=eps, dtype=dtype
             )
         else:
-                self.norm = LayerNorm(
+            raise NotImplementedError(f"Norm type {self.norm_type} not implemented")
-                    hidden_size,
-                    elementwise_affine=elementwise_affine,
-                    eps=eps,
-                    dtype=dtype,
-                )
-        else:
-            raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
-    def forward(
+    def forward_cuda(
         self,
         residual: torch.Tensor,
         x: torch.Tensor,
@@ -306,18 +306,45 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
         shift: torch.Tensor,
         scale: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
+        if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192:
-        Apply gated residual connection, followed by layernorm and
+            import warnings
-        scale/shift in a single fused operation.
 
-        Returns:
+            warnings.warn(
-            Tuple containing:
+                "FusedScaleResidualNormScaleShift cuda not available, using native fallback",
-            - normalized and modulated output of shape: [batch_size, seq_len, inner_dim]
+                stacklevel=2,
-            - residual value (value after residual connection
+            )
-              but before normalization)
+            return self.forward_native(residual, x, gate, shift, scale)
-        """
+
+        from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import (
+            fused_scale_residual_norm_scale_shift,
+        )
+
+        return fused_scale_residual_norm_scale_shift(
+            residual.contiguous(),
+            x.contiguous(),
+            gate.contiguous() if isinstance(gate, torch.Tensor) else None,
+            _ensure_contiguous(getattr(self.norm, "weight", None)),
+            _ensure_contiguous(getattr(self.norm, "bias", None)),
+            scale.contiguous(),
+            shift.contiguous(),
+            self.norm_type,
+            self.eps,
+        )
+
+    def forward_hip(self, *args, **kwargs):
+        # ROCm does not support CUDA/CUTLASS-based fused kernels yet,
+        # so we fall back to the native PyTorch implementation.
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(
+        self,
+        residual: torch.Tensor,
+        x: torch.Tensor,
+        gate: torch.Tensor | int,
+        shift: torch.Tensor,
+        scale: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # x.shape: [batch_size, seq_len, inner_dim]
-        # Apply residual connection with gating
         if isinstance(gate, int):
             # used by cross-attention, should be 1
             assert gate == 1
@@ -331,91 +358,97 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
                     x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate
                 ).flatten(1, 2)
             else:
-                # used by bidirectional self attention
                 # gate.shape: [batch_size, 1, inner_dim]
                 residual_output = residual + x * gate
         else:
             raise ValueError(f"Gate type {type(gate)} not supported")
-        # residual_output.shape: [batch_size, seq_len, inner_dim]
-
-        # Apply normalization
         normalized = self.norm(residual_output)
-
+        modulated = fuse_scale_shift_kernel(normalized, scale, shift)
-        # modulated = fused_scale_shift(
-        #     normalized,
-        #     scale,
-        #     shift,
-        # )
-        modulated = fuse_scale_shift_kernel(
-            normalized,
-            scale,
-            shift,
-        )
         return modulated, residual_output
 
 
-class LayerNormScaleShift(nn.Module):
+class ScaleResidualLayerNormScaleShift(_ScaleResidualNormScaleShift):
+    norm_type = "layer"
+
+
+class ScaleResidualRMSNormScaleShift(_ScaleResidualNormScaleShift):
+    norm_type = "rms"
+
+
+class _NormScaleShift(CustomOp):
     """
-    Fused operation that combines LayerNorm with scale and shift operations.
+    Fused kernel that combines:
-    This reduces memory bandwidth by combining memory-bound operations.
+    1. normed = layernorm(x) or rmsnorm(x)
+    2. out = normed * (1 + scale) + shift
+    compute_dtype is always fp32 for higher precision.
     """
 
+    norm_type: str
+
     def __init__(
         self,
         hidden_size: int,
-        norm_type: str = "rms",
         eps: float = 1e-6,
         elementwise_affine: bool = False,
         dtype: torch.dtype = torch.float32,
-        compute_dtype: torch.dtype | None = None,
         prefix: str = "",
     ):
         super().__init__()
-        self.compute_dtype = compute_dtype
+        self.eps = eps
-        if norm_type == "rms":
+        if self.norm_type == "rms":
-            self.norm = RMSNorm(hidden_size, has_weight=elementwise_affine, eps=eps)
+            self.norm = RMSNorm(hidden_size, eps=eps, dtype=dtype)
-        elif norm_type == "layer":
+        elif self.norm_type == "layer":
-            if self.compute_dtype == torch.float32:
             self.norm = FP32LayerNorm(
-                    hidden_size, elementwise_affine=elementwise_affine, eps=eps
+                hidden_size, elementwise_affine=elementwise_affine, eps=eps, dtype=dtype
             )
         else:
-                self.norm = nn.LayerNorm(
+            raise NotImplementedError(f"Norm type {self.norm_type} not implemented")
-                    hidden_size,
-                    elementwise_affine=elementwise_affine,
-                    eps=eps,
-                    dtype=dtype,
-                )
-        else:
-            raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
-    def forward(
+    def forward_cuda(
+        self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
+    ) -> torch.Tensor:
+        if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192:
+            import warnings
+
+            warnings.warn(
+                "FusedNormScaleShift cuda not available, using native fallback",
+                stacklevel=2,
+            )
+            return self.forward_native(x, shift, scale)
+
+        from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import (
+            fused_norm_scale_shift,
+        )
+
+        return fused_norm_scale_shift(
+            x.contiguous(),
+            _ensure_contiguous(getattr(self.norm, "weight", None)),
+            _ensure_contiguous(getattr(self.norm, "bias", None)),
+            scale.contiguous(),
+            shift.contiguous(),
+            self.norm_type,
+            self.eps,
+        )
+
+    def forward_hip(self, *args, **kwargs):
+        # ROCm does not support CUDA/CUTLASS-based fused kernels yet,
+        # so we fall back to the native PyTorch implementation.
+        return self.forward_native(*args, **kwargs)
+
+    def forward_native(
         self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
     ) -> torch.Tensor:
-        """Apply ln followed by scale and shift in a single fused operation."""
-        # x.shape: [batch_size, seq_len, inner_dim]
         normalized = self.norm(x)
-        if self.compute_dtype == torch.float32:
+        modulated = fuse_scale_shift_kernel(normalized, scale, shift)
-            normalized = normalized.float()
+        return modulated.to(x.dtype)
 
-        if scale.dim() == 4:
-            # scale.shape: [batch_size, num_frames, 1, inner_dim]
-            num_frames = scale.shape[1]
-            frame_seqlen = normalized.shape[1] // num_frames
-            output = (
-                normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen))
-                * (1.0 + scale)
-                + shift
-            ).flatten(1, 2)
-        else:
-            # scale.shape: [batch_size, 1, inner_dim]
-            # shift.shape: [batch_size, 1, inner_dim]
-            output = normalized * (1.0 + scale) + shift
 
-        if self.compute_dtype == torch.float32:
+class LayerNormScaleShift(_NormScaleShift):
-            output = output.to(x.dtype)
+    norm_type = "layer"
 
-        return output
+
+class RMSNormScaleShift(_NormScaleShift):
+    norm_type = "rms"
 
 
 def apply_qk_norm(
@@ -470,3 +503,14 @@ def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor:
     )
     output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight
     return output.to(dtype=src_dtype)
+
+
+# TODO: Workaround, fuse norm with new select01 kernel
+def apply_layernorm_only(x: torch.Tensor, layernorm_scale_shift: LayerNormScaleShift):
+    return norm_infer(
+        x.view(-1, x.shape[-1]),
+        layernorm_scale_shift.norm.weight,
+        layernorm_scale_shift.norm.bias,
+        eps=layernorm_scale_shift.eps,
+        is_rms_norm=False,
+    ).view(x.shape)

python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py

@@ -296,24 +296,14 @@ class CausalWanTransformerBlock(nn.Module):
             raise Exception
         assert cross_attn_norm is True
         self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
-            dim,
+            dim, eps=eps, elementwise_affine=True, dtype=torch.float32
-            norm_type="layer",
-            eps=eps,
-            elementwise_affine=True,
-            dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
 
         # 2. Cross-attention
         # Only T2V for now
         self.attn2 = WanT2VCrossAttention(dim, num_heads, qk_norm=qk_norm, eps=eps)
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
-            dim,
+            dim, eps=eps, elementwise_affine=False, dtype=torch.float32
-            norm_type="layer",
-            eps=eps,
-            elementwise_affine=False,
-            dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
 
         # 3. Feed-forward
@@ -484,11 +474,9 @@ class CausalWanTransformer3DModel(BaseDiT, OffloadableDiTMixin):
         # 4. Output norm & projection
         self.norm_out = LayerNormScaleShift(
             inner_dim,
-            norm_type="layer",
             eps=config.eps,
             elementwise_affine=False,
             dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
         self.proj_out = nn.Linear(
             inner_dim, config.out_channels * math.prod(config.patch_size)

python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py

@@ -76,10 +76,10 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Fused operations for image stream
         self.img_attn_norm = LayerNormScaleShift(
-            hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype
+            hidden_size, elementwise_affine=False, dtype=dtype
         )
         self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
-            hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype
+            hidden_size, elementwise_affine=False, dtype=dtype
         )
         self.img_mlp_residual = MulAdd()
 
@@ -122,10 +122,10 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Fused operations for text stream
         self.txt_attn_norm = LayerNormScaleShift(
-            hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype
+            hidden_size, elementwise_affine=False, dtype=dtype
         )
         self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
-            hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype
+            hidden_size, elementwise_affine=False, dtype=dtype
         )
         self.txt_mlp_residual = MulAdd()
 
@@ -299,7 +299,6 @@ class MMSingleStreamBlock(nn.Module):
         # Fused operations with better naming
         self.input_norm_scale_shift = LayerNormScaleShift(
             hidden_size,
-            norm_type="layer",
             eps=1e-6,
             elementwise_affine=False,
             dtype=dtype,

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -19,8 +19,10 @@ from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
 from sglang.multimodal_gen.runtime.layers.layernorm import (
-    LayerNorm,
+    LayerNormScaleShift,
     RMSNorm,
+    ScaleResidualLayerNormScaleShift,
+    apply_layernorm_only,
     apply_qk_norm,
 )
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
@@ -646,7 +648,9 @@ class QwenImageTransformerBlock(nn.Module):
                 dim, 6 * dim, bias=True
             ),  # For scale, shift, gate for norm1 and norm2
         )
-        self.img_norm1 = LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_norm1 = LayerNormScaleShift(
+            hidden_size=dim, eps=eps, elementwise_affine=False
+        )
 
         self.attn = QwenImageCrossAttention(
             dim=dim,
@@ -655,7 +659,9 @@ class QwenImageTransformerBlock(nn.Module):
             context_pre_only=False,
             head_dim=attention_head_dim,
         )
-        self.img_norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.img_norm2 = ScaleResidualLayerNormScaleShift(
+            hidden_size=dim, eps=eps, elementwise_affine=False
+        )
         self.img_mlp = FeedForward(
             dim=dim, dim_out=dim, activation_fn="gelu-approximate"
         )
@@ -667,16 +673,37 @@ class QwenImageTransformerBlock(nn.Module):
                 dim, 6 * dim, bias=True
             ),  # For scale, shift, gate for norm1 and norm2
         )
-        self.txt_norm1 = LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_norm1 = LayerNormScaleShift(
+            hidden_size=dim, eps=eps, elementwise_affine=False
+        )
         # Text doesn't need separate attention - it's handled by img_attn joint computation
-        self.txt_norm2 = LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_norm2 = ScaleResidualLayerNormScaleShift(
+            hidden_size=dim, eps=eps, elementwise_affine=False
+        )
         self.txt_mlp = FeedForward(
             dim=dim, dim_out=dim, activation_fn="gelu-approximate"
         )
         # Utils
         self.fuse_mul_add = MulAdd()
 
-    def _modulate(self, x, mod_params, index=None):
+    def _modulate(
+        self,
+        x: torch.Tensor,
+        mod_params: torch.Tensor,
+        norm_module: Union[LayerNormScaleShift, ScaleResidualLayerNormScaleShift],
+        index: Optional[torch.Tensor] = None,
+        gate_x: Optional[torch.Tensor] = None,
+        residual_x: Optional[torch.Tensor] = None,
+    ) -> Union[
+        Tuple[torch.Tensor, torch.Tensor],
+        Tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+    ]:
+        # Apply attention gates and add residual (like in Megatron)
+        #   - residual_out = gate_x * x + residual_x
+        # - x = norm(residual_out) * (1 + scale) + shift
+        # TODO: clean code here
+        is_scale_residual = isinstance(norm_module, ScaleResidualLayerNormScaleShift)
+
         shift, scale, gate = mod_params.chunk(3, dim=-1)
         if index is not None:
             actual_batch = x.shape[0]
@@ -689,12 +716,16 @@ class QwenImageTransformerBlock(nn.Module):
                 scale[actual_batch : 2 * actual_batch],
             )
             gate0, gate1 = gate[:actual_batch], gate[actual_batch : 2 * actual_batch]
-
             if _is_cuda:
+                if is_scale_residual:
+                    x = gate_x * x + residual_x
+                    residual_out = x
                 if not x.is_contiguous():
                     x = x.contiguous()
                 if not index.is_contiguous():
                     index = index.contiguous()
+                # TODO: fuse norm with above select01 kernel, workaround now
+                x = apply_layernorm_only(x, norm_module)
                 x, gate_result = fuse_scale_shift_gate_select01_kernel(
                     x,
                     scale0=scale0.contiguous(),
@@ -705,6 +736,9 @@ class QwenImageTransformerBlock(nn.Module):
                     gate1=gate1.contiguous(),
                     index=index,
                 )
+                if is_scale_residual:
+                    return x, residual_out, gate_result
+                else:
                     return x, gate_result
             else:
                 mask = (index == 0).unsqueeze(-1)
@@ -715,15 +749,34 @@ class QwenImageTransformerBlock(nn.Module):
                     mask, scale0.unsqueeze(1), scale1.unsqueeze(1)
                 )
                 gate_result = torch.where(mask, gate0.unsqueeze(1), gate1.unsqueeze(1))
-                return (
+                if is_scale_residual:
-                    self.fuse_mul_add(x, scale_result, shift_result, k=1.0),
+                    modulated, residual_out = norm_module(
-                    gate_result,
+                        residual=residual_x,
+                        x=x,
+                        gate=gate_x,
+                        shift=shift_result,
+                        scale=scale_result,
                     )
+                    return modulated, residual_out, gate_result
+                else:
+                    modulated = norm_module(x=x, shift=shift_result, scale=scale_result)
+                    return modulated, gate_result
         else:
             shift_result = shift.unsqueeze(1)
             scale_result = scale.unsqueeze(1)
             gate_result = gate.unsqueeze(1)
-            return self.fuse_mul_add(x, scale_result, shift_result, k=1.0), gate_result
+            if is_scale_residual:
+                modulated, residual_out = norm_module(
+                    residual=residual_x,
+                    x=x,
+                    gate=gate_x,
+                    shift=shift_result,
+                    scale=scale_result,
+                )
+                return modulated, residual_out, gate_result
+            else:
+                modulated = norm_module(x=x, shift=shift_result, scale=scale_result)
+                return modulated, gate_result
 
     def forward(
         self,
@@ -745,13 +798,15 @@ class QwenImageTransformerBlock(nn.Module):
         txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
 
         # Process image stream - norm1 + modulation
-
+        img_modulated, img_gate1 = self._modulate(
-        img_normed = self.img_norm1(hidden_states)
+            hidden_states, img_mod1, self.img_norm1, modulate_index
-
+        )
-        img_modulated, img_gate1 = self._modulate(img_normed, img_mod1, modulate_index)
         # Process text stream - norm1 + modulation
-        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_shift1, txt_scale1, txt_gate1_raw = txt_mod1.chunk(3, dim=-1)
-        txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
+        txt_modulated = self.txt_norm1(
+            encoder_hidden_states, shift=txt_shift1, scale=txt_scale1
+        )
+        txt_gate1 = txt_gate1_raw.unsqueeze(1)
 
         # Use QwenAttnProcessor2_0 for joint attention computation
         # This directly implements the DoubleStreamLayerMegatron logic:
@@ -772,23 +827,28 @@ class QwenImageTransformerBlock(nn.Module):
 
         # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
         img_attn_output, txt_attn_output = attn_output
-
-        # Apply attention gates and add residual (like in Megatron)
-        hidden_states = hidden_states + img_gate1 * img_attn_output
-
-        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
-
         # Process image stream - norm2 + MLP
-        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, hidden_states, img_gate2 = self._modulate(
-        img_modulated2, img_gate2 = self._modulate(
+            img_attn_output,
-            img_normed2, img_mod2, modulate_index
+            img_mod2,
+            self.img_norm2,
+            modulate_index,
+            gate_x=img_gate1,
+            residual_x=hidden_states,
         )
         img_mlp_output = self.img_mlp(img_modulated2)
         hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states)
 
         # Process text stream - norm2 + MLP
-        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1)
-        txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
+        txt_modulated2, encoder_hidden_states = self.txt_norm2(
+            residual=encoder_hidden_states,
+            x=txt_attn_output,
+            gate=txt_gate1,
+            shift=txt_shift2,
+            scale=txt_scale2,
+        )
+        txt_gate2 = txt_gate2_raw.unsqueeze(1)
         txt_mlp_output = self.txt_mlp(txt_modulated2)
         encoder_hidden_states = self.fuse_mul_add(
             txt_mlp_output, txt_gate2, encoder_hidden_states

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -298,7 +298,12 @@ class WanTransformerBlock(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = LayerNormScaleShift(
+            dim,
+            eps=eps,
+            elementwise_affine=False,
+            dtype=torch.float32,
+        )
         self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
         self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
         self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
@@ -344,11 +349,9 @@ class WanTransformerBlock(nn.Module):
         self.tp_rmsnorm = qk_norm == "rms_norm_across_heads" and tp_size > 1
         self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
-            norm_type="layer",
             eps=eps,
             elementwise_affine=True,
             dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
 
         # 2. Cross-attention
@@ -372,11 +375,9 @@ class WanTransformerBlock(nn.Module):
             )
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
-            norm_type="layer",
             eps=eps,
             elementwise_affine=False,
             dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
 
         # 3. Feed-forward
@@ -418,8 +419,7 @@ class WanTransformerBlock(nn.Module):
         assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        norm1 = self.norm1(hidden_states.float())
+        norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa)
-        norm_hidden_states = (norm1 * (1 + scale_msa) + shift_msa).to(orig_dtype)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
@@ -506,7 +506,12 @@ class WanTransformerBlock_VSA(nn.Module):
         super().__init__()
 
         # 1. Self-attention
-        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.norm1 = LayerNormScaleShift(
+            dim,
+            eps=eps,
+            elementwise_affine=False,
+            dtype=torch.float32,
+        )
         self.to_q = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
         self.to_k = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
         self.to_v = ColumnParallelLinear(dim, dim, bias=True, gather_output=True)
@@ -538,11 +543,9 @@ class WanTransformerBlock_VSA(nn.Module):
         assert cross_attn_norm is True
         self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
-            norm_type="layer",
             eps=eps,
             elementwise_affine=True,
             dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
 
         if AttentionBackendEnum.VIDEO_SPARSE_ATTN in supported_attention_backends:
@@ -568,11 +571,9 @@ class WanTransformerBlock_VSA(nn.Module):
             )
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
-            norm_type="layer",
             eps=eps,
             elementwise_affine=False,
             dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
 
         # 3. Feed-forward
@@ -600,9 +601,7 @@ class WanTransformerBlock_VSA(nn.Module):
         assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
-        norm_hidden_states = (
+        norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa)
-            self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa
-        ).to(orig_dtype)
         query, _ = self.to_q(norm_hidden_states)
         key, _ = self.to_k(norm_hidden_states)
         value, _ = self.to_v(norm_hidden_states)
@@ -736,11 +735,9 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin):
         # 4. Output norm & projection
         self.norm_out = LayerNormScaleShift(
             inner_dim,
-            norm_type="layer",
             eps=config.eps,
             elementwise_affine=False,
             dtype=torch.float32,
-            compute_dtype=torch.float32,
         )
         self.proj_out = nn.Linear(
             inner_dim, config.out_channels * math.prod(config.patch_size)