diff --git a/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py b/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py new file mode 100644 index 000000000..e1c7d72af --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/common/norm_fusion.py b/python/sglang/jit_kernel/diffusion/cutedsl/common/norm_fusion.py new file mode 100644 index 000000000..01b080284 --- /dev/null +++ b/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 diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py b/python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py new file mode 100644 index 000000000..246bc1856 --- /dev/null +++ b/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 diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py b/python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py new file mode 100644 index 000000000..a786dba6e --- /dev/null +++ b/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 = * 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 + + 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) * + + 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 + ) + + 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"') diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/utils.py b/python/sglang/jit_kernel/diffusion/cutedsl/utils.py new file mode 100644 index 000000000..d23c2342b --- /dev/null +++ b/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, +} diff --git a/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py b/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py new file mode 100644 index 000000000..e7c45041e --- /dev/null +++ b/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__]) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 03989a55e..eaf87151f 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/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.bias.float() if self.bias is not None else None, + self.weight.float().to(device=device) if self.weight 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 operation that combines: - 1. Gated residual connection - 2. LayerNorm - 3. Scale and shift operations +################################################################################ +# Fused norm kernel +################################################################################ +def _ensure_contiguous(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + return tensor.contiguous() if tensor is not None else None - 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.norm = RMSNorm( - hidden_size, has_weight=elementwise_affine, eps=eps, dtype=dtype + self.eps = eps + self.dtype = dtype + if self.norm_type == "rms": + self.norm = RMSNorm(hidden_size, eps=eps, dtype=dtype) + elif self.norm_type == "layer": + self.norm = FP32LayerNorm( + hidden_size, elementwise_affine=elementwise_affine, eps=eps, dtype=dtype ) - elif norm_type == "layer": - if compute_dtype == torch.float32: - self.norm = FP32LayerNorm( - hidden_size, elementwise_affine=elementwise_affine, eps=eps - ) - else: - self.norm = LayerNorm( - hidden_size, - elementwise_affine=elementwise_affine, - eps=eps, - dtype=dtype, - ) else: - raise NotImplementedError(f"Norm type {norm_type} not implemented") + raise NotImplementedError(f"Norm type {self.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]: - """ - Apply gated residual connection, followed by layernorm and - scale/shift in a single fused operation. + if x.shape[-1] % 256 != 0 and x.shape[-1] <= 8192: + import warnings - Returns: - Tuple containing: - - normalized and modulated output of shape: [batch_size, seq_len, inner_dim] - - residual value (value after residual connection - but before normalization) - """ + warnings.warn( + "FusedScaleResidualNormScaleShift cuda not available, using native fallback", + stacklevel=2, + ) + 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 = fused_scale_shift( - # normalized, - # scale, - # shift, - # ) - modulated = fuse_scale_shift_kernel( - 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. - This reduces memory bandwidth by combining memory-bound operations. + Fused kernel that combines: + 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 - if norm_type == "rms": - self.norm = RMSNorm(hidden_size, has_weight=elementwise_affine, eps=eps) - elif norm_type == "layer": - if self.compute_dtype == torch.float32: - self.norm = FP32LayerNorm( - hidden_size, elementwise_affine=elementwise_affine, eps=eps - ) - else: - self.norm = nn.LayerNorm( - hidden_size, - elementwise_affine=elementwise_affine, - eps=eps, - dtype=dtype, - ) + self.eps = eps + if self.norm_type == "rms": + self.norm = RMSNorm(hidden_size, eps=eps, dtype=dtype) + elif self.norm_type == "layer": + self.norm = FP32LayerNorm( + hidden_size, elementwise_affine=elementwise_affine, eps=eps, dtype=dtype + ) else: - raise NotImplementedError(f"Norm type {norm_type} not implemented") + raise NotImplementedError(f"Norm type {self.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: - normalized = normalized.float() + modulated = fuse_scale_shift_kernel(normalized, scale, shift) + 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: - output = output.to(x.dtype) +class LayerNormScaleShift(_NormScaleShift): + 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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py index 0095a0591..d2e5c7ed4 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py +++ b/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, - norm_type="layer", - eps=eps, - elementwise_affine=True, - dtype=torch.float32, - compute_dtype=torch.float32, + dim, eps=eps, elementwise_affine=True, 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, - norm_type="layer", - eps=eps, - elementwise_affine=False, - dtype=torch.float32, - compute_dtype=torch.float32, + dim, eps=eps, elementwise_affine=False, 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) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py index 80539de66..8ef9162d1 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index ad5c96229..2d502bc7a 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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,7 +736,10 @@ class QwenImageTransformerBlock(nn.Module): gate1=gate1.contiguous(), index=index, ) - return x, gate_result + if is_scale_residual: + return x, residual_out, gate_result + else: + return x, gate_result else: mask = (index == 0).unsqueeze(-1) shift_result = torch.where( @@ -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 ( - 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 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_normed = self.img_norm1(hidden_states) - - img_modulated, img_gate1 = self._modulate(img_normed, img_mod1, modulate_index) + img_modulated, img_gate1 = self._modulate( + hidden_states, img_mod1, self.img_norm1, modulate_index + ) # Process text stream - norm1 + modulation - txt_normed = self.txt_norm1(encoder_hidden_states) - txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1) + txt_shift1, txt_scale1, txt_gate1_raw = txt_mod1.chunk(3, dim=-1) + 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, img_gate2 = self._modulate( - img_normed2, img_mod2, modulate_index + img_modulated2, hidden_states, img_gate2 = self._modulate( + img_attn_output, + 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_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2) + txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1) + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index ed980d079..d80a38fec 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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 = (norm1 * (1 + scale_msa) + shift_msa).to(orig_dtype) + norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) 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 = ( - self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa - ).to(orig_dtype) + norm_hidden_states = self.norm1(hidden_states, shift_msa, scale_msa) 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)