[diffusion] kernel: gated residual layernorm scale shift and layernorm scale shift kernel fusion for Qwen-Image, WAN and HunyuanVideo (#14717)
Co-authored-by: AichenF <aichenf@nvidia.com> Co-authored-by: jianyingzhu <joeyzhu@nvidia.com> Co-authored-by: root <root@a4u8g-0120.ipp2a2.colossus.nvidia.com> Co-authored-by: Yihan Chen <yingluosanqian@example.com> Co-authored-by: 陈一涵 <yingluosanqian@gmail.com> Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
Jianying
201
python/sglang/jit_kernel/diffusion/cutedsl/common/norm_fusion.py
Normal file
201
python/sglang/jit_kernel/diffusion/cutedsl/common/norm_fusion.py
Normal 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
33
python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py
Normal 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
|
||||
@@ -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
10
python/sglang/jit_kernel/diffusion/cutedsl/utils.py
Normal 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,
|
||||
}
|
||||
Reference in New Issue
Block a user