chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Jianying
2026-02-04 13:46:20 +08:00
committed by GitHub
parent 669a9bd180
commit 4739f2e8d5
11 changed files with 1285 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files

134
python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py Normal file
View File

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

201
python/sglang/jit_kernel/diffusion/cutedsl/common/norm_fusion.py Normal file
View File

@@ -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

33
python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py Normal file
View File

@@ -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

419
python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py Normal file
View File

@@ -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"')

10
python/sglang/jit_kernel/diffusion/cutedsl/utils.py Normal file
View File

@@ -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,
}

238
python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py Normal file
View File

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