[diffusion] improve: norm fusion for z-image (#18762)
Signed-off-by: Chi McIsaac <chixie.mcisaac@gmail.com> Co-authored-by: yihanc <yingluosanqian@gmail.com> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Chi McIsaac
co-authored by
yihanc
Mick
yihanc
Mick
parent
ef13031243
commit
005e582d06
@@ -0,0 +1,379 @@
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import cuda.bindings.driver as cuda
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.diffusion.cutedsl.common.norm_fusion import (
|
||||
apply_norm_cta,
|
||||
broadcast_tensor_for_bsfd,
|
||||
tensor_slice_for_bsfd,
|
||||
)
|
||||
from sglang.jit_kernel.diffusion.cutedsl.utils import 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 NormTanhMulAddNormScale:
|
||||
@classmethod
|
||||
def make_hash_key(cls, *inputs):
|
||||
"""
|
||||
Compile-time values:
|
||||
- D: hidden dimension (size of the last dimension)
|
||||
- norm_type: layer norm or RMS norm
|
||||
- tensor dtype
|
||||
- tensor rank (i.e., tensor.ndim)
|
||||
|
||||
Runtime values:
|
||||
- all other inputs
|
||||
|
||||
This hash key defines the compile-time specialization boundary for
|
||||
NormTanhMulAddNormScale kernels.
|
||||
"""
|
||||
|
||||
def _sig(val):
|
||||
if isinstance(val, torch.Tensor):
|
||||
return (val.dtype, val.ndim, val.shape[-1])
|
||||
return val
|
||||
|
||||
return tuple(_sig(val) for val in inputs)
|
||||
|
||||
def __init__(self, D: int, norm_type: str, is_norm2: bool):
|
||||
self.D = D
|
||||
self.norm_type = norm_type # "layer" or "rms"
|
||||
self.is_norm2 = is_norm2 # single norm or double norm
|
||||
self.num_warps = self.D // 256 # num of warps per cta
|
||||
self.num_threads = self.num_warps * WARP_SIZE # num of threads per cta
|
||||
|
||||
@cute.jit
|
||||
def __call__(
|
||||
self,
|
||||
mY,
|
||||
mY2,
|
||||
mX,
|
||||
mWeight,
|
||||
mBias,
|
||||
mScale,
|
||||
mShift,
|
||||
mWeight2,
|
||||
mBias2,
|
||||
mScale2,
|
||||
eps: cutlass.Float32 = cutlass.Float32(1e-5),
|
||||
stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
|
||||
):
|
||||
# Tensor shapes
|
||||
B, S, _ = mX.shape # (batch, seq_len, hidden_dim)
|
||||
# Vectorized copy configuration
|
||||
num_vectorized = 8 # maximum num of elem per copy
|
||||
atom_copy = cute.make_copy_atom(
|
||||
cute.nvgpu.CopyUniversalOp(),
|
||||
mX.element_type,
|
||||
num_bits_per_copy=128,
|
||||
)
|
||||
# Thread/value layouts for tiled copy
|
||||
t_layout = cute.make_layout(self.num_threads) # thread layout within a CTA
|
||||
v_layout = cute.make_layout(num_vectorized) # per-thread vector layout
|
||||
tiled_copy = cute.make_tiled_copy_tv(atom_copy, t_layout, v_layout)
|
||||
|
||||
self.kernel(
|
||||
mY,
|
||||
mY2,
|
||||
mX,
|
||||
mWeight,
|
||||
mBias,
|
||||
mScale,
|
||||
mShift,
|
||||
mWeight2,
|
||||
mBias2,
|
||||
mScale2,
|
||||
tiled_copy,
|
||||
eps,
|
||||
).launch(
|
||||
grid=[B * S, 1, 1],
|
||||
block=[self.num_threads, 1, 1],
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
@cute.kernel
|
||||
def kernel(
|
||||
self,
|
||||
mY,
|
||||
mY2,
|
||||
mX,
|
||||
mWeight,
|
||||
mBias,
|
||||
mScale,
|
||||
mShift,
|
||||
mWeight2,
|
||||
mBias2,
|
||||
mScale2,
|
||||
tiled_copy: cute.TiledCopy,
|
||||
eps: cutlass.Float32,
|
||||
):
|
||||
_, S, _ = mX.shape
|
||||
tidx, _, _ = cute.arch.thread_idx() # thread index
|
||||
bid, _, _ = cute.arch.block_idx() # cta index
|
||||
bidx = cutlass.Int32(bid // S) # batch index
|
||||
bidy = cutlass.Int32(bid % S) # seq_len index
|
||||
thr_copy = tiled_copy.get_slice(tidx)
|
||||
|
||||
@cute.jit
|
||||
def slice_if(mV):
|
||||
if cutlass.const_expr(isinstance(mV, cute.Tensor)):
|
||||
return tensor_slice_for_bsfd(mV, thr_copy, bidx, bidy, S, self.D)
|
||||
return mV, mV
|
||||
|
||||
@cute.jit
|
||||
def copy_if(src, dst):
|
||||
if cutlass.const_expr(
|
||||
isinstance(src, cute.Tensor) and isinstance(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)
|
||||
tXgX, tXrX = slice_if(mX) # x
|
||||
tWgW, tWrW = slice_if(mWeight) # weight
|
||||
tBgB, tBrB = slice_if(mBias) # bias
|
||||
tSCgSC, tSCrSC = slice_if(mScale) # scale
|
||||
tSHgSH, tSHrSH = slice_if(mShift) # shift
|
||||
tYgY, tYrY = slice_if(mY) # y
|
||||
if cutlass.const_expr(self.is_norm2):
|
||||
tYgY2, tYrY2 = slice_if(mY2) # y2
|
||||
tWgW2, tWrW2 = slice_if(mWeight2) # weight2
|
||||
tBgB2, tBrB2 = slice_if(mBias2) # bias2
|
||||
tSCgSC2, tSCrSC2 = slice_if(mScale2) # scale2
|
||||
# Load: load tensor from global memory to registers
|
||||
copy_if(tXgX, tXrX) # gmem -> rmem
|
||||
copy_if(tWgW, tWrW) # gmem -> rmem
|
||||
copy_if(tBgB, tBrB) # gmem -> rmem
|
||||
tNrN = norm(tXrX, tWrW, tBrB)
|
||||
# Compute: value = value * tanh(<scale>) + <shift>
|
||||
copy_if(tSCgSC, tSCrSC) # gmem -> rmem
|
||||
copy_if(tSHgSH, tSHrSH) # gmem -> rmem
|
||||
value = tNrN.load() * cute.tanh(tSCrSC.load()) + tSHrSH.load()
|
||||
# Store: y
|
||||
tYrY.store(value.to(tYrY.element_type))
|
||||
copy_if(tYrY, tYgY) # rmem -> gmem
|
||||
if cutlass.const_expr(self.is_norm2):
|
||||
copy_if(tWgW2, tWrW2) # gmem -> rmem
|
||||
copy_if(tBgB2, tBrB2) # gmem -> rmem
|
||||
tNrN2 = norm(tYrY, tWrW2, tBrB2)
|
||||
# Compute: value2 = value2 * (1 + <scale2>)
|
||||
copy_if(tSCgSC2, tSCrSC2) # gmem -> rmem
|
||||
value2 = tNrN2.load() * (1 + tSCrSC2.load())
|
||||
# Store: y2
|
||||
tYrY2.store(value2.to(tYrY2.element_type))
|
||||
copy_if(tYrY2, tYgY2) # rmem -> gmem
|
||||
|
||||
|
||||
def validate_3d(t: torch.Tensor, B: int, S: int, D: int):
|
||||
if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
|
||||
if (
|
||||
t.ndim != 3
|
||||
or (t.shape[0] not in (1, B))
|
||||
or (t.shape[1] not in (1, S) or t.shape[2] != D)
|
||||
):
|
||||
raise ValueError(f"Validate failed: unsupported 3d-tensor: {t.shape}.")
|
||||
if t.stride()[-1] != 1:
|
||||
raise ValueError(f"Validate failed: not contiguous on dim D.")
|
||||
|
||||
|
||||
def validate_weight_bias(t: Optional[torch.Tensor], D: int):
|
||||
if t is None:
|
||||
return
|
||||
if t.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError(f"Validate failed: unsupported dtype: {t.dtype}")
|
||||
if t.shape != (D,):
|
||||
raise ValueError(f"Validate failed: unsupported tensor shape: {t.shape}.")
|
||||
if t.stride()[-1] != 1:
|
||||
raise ValueError(f"Validate failed: not contiguous on dim D.")
|
||||
|
||||
|
||||
@torch.library.custom_op("sglang::fused_norm_tanh_mul_add", mutates_args=())
|
||||
def fused_norm_tanh_mul_add(
|
||||
x: torch.Tensor,
|
||||
weight: Optional[torch.Tensor],
|
||||
bias: Optional[torch.Tensor],
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
norm_type: str,
|
||||
eps: float = 1e-5,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Fuse: norm(x) * tanh(scale) + shift
|
||||
where norm is either layernorm or rmsnorm.
|
||||
|
||||
Expects:
|
||||
- x: [B, S, D]
|
||||
- weight/bias: None, [D]
|
||||
- scale/shift: [1/B, 1/S, D]
|
||||
- norm_type: str, "layer" or "rms"
|
||||
- eps: Optional[float], default: 1e-5
|
||||
|
||||
D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
|
||||
thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
|
||||
"""
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
# Tensor Validation
|
||||
BSD = x.shape
|
||||
validate_3d(x, *BSD)
|
||||
validate_weight_bias(weight, BSD[2])
|
||||
validate_weight_bias(bias, BSD[2])
|
||||
validate_3d(scale, *BSD)
|
||||
validate_3d(shift, *BSD)
|
||||
if norm_type == "layer" or norm_type == "rms":
|
||||
D = x.shape[-1]
|
||||
if D % 256 != 0 or D > 8192:
|
||||
raise ValueError(
|
||||
f"D={D} not supported, must be multiple of 256 and <= 8192"
|
||||
)
|
||||
y = torch.empty_like(x) # create output tensor
|
||||
scale = broadcast_tensor_for_bsfd(scale, *x.shape) # handle various shapes
|
||||
shift = broadcast_tensor_for_bsfd(shift, *x.shape) # handle various shapes
|
||||
# y2, weight2, bias2, scale2 is None
|
||||
torch_tensors = [y, None, x, weight, bias, scale, shift, None, None, None]
|
||||
cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
|
||||
# Compile cache
|
||||
hash_key = NormTanhMulAddNormScale.make_hash_key(norm_type, *torch_tensors)
|
||||
compiled_fn = _COMPILE_CACHE.get(hash_key)
|
||||
if compiled_fn is None:
|
||||
kernel = NormTanhMulAddNormScale(D, norm_type, is_norm2=False)
|
||||
fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
|
||||
compiled_fn = cute.compile(
|
||||
kernel, *fake_sig_args, options="--enable-tvm-ffi"
|
||||
)
|
||||
_COMPILE_CACHE[hash_key] = compiled_fn
|
||||
# Execute
|
||||
compiled_fn(*cute_tensor_args, eps, stream)
|
||||
return y
|
||||
else:
|
||||
raise ValueError(f'norm_type must be one of "layer" and "rms"')
|
||||
|
||||
|
||||
@fused_norm_tanh_mul_add.register_fake
|
||||
def _fused_norm_tanh_mul_add_fake(x, weight, bias, scale, shift, norm_type, eps=1e-5):
|
||||
return x.new_empty(x.shape)
|
||||
|
||||
|
||||
@torch.library.custom_op("sglang::fused_norm_tanh_mul_add_norm_scale", mutates_args=())
|
||||
def fused_norm_tanh_mul_add_norm_scale(
|
||||
x: torch.Tensor,
|
||||
weight: Optional[torch.Tensor],
|
||||
bias: Optional[torch.Tensor],
|
||||
scale: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
weight2: Optional[torch.Tensor],
|
||||
bias2: Optional[torch.Tensor],
|
||||
scale2: torch.Tensor,
|
||||
norm_type: str,
|
||||
eps: float = 1e-5,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Fuse:
|
||||
y = norm(x) * tanh(scale) + shift
|
||||
y2 = norm(y) * (1 + scale2)
|
||||
where norm is either layernorm or rmsnorm.
|
||||
|
||||
Expects:
|
||||
- x: [B, S, D]
|
||||
- weight/bia/weight2/bias2: None, [D]
|
||||
- scale/shift/scale2: [1/B, 1/S, D]
|
||||
- norm_type: str, "layer" or "rms"
|
||||
- eps: Optional[float], default: 1e-5
|
||||
|
||||
D must be a multiple of 256 and <= 8192 to enable LDG.128 vectorized loads per
|
||||
thread and avoid predicated loads (e.g., bounds checks such as `index < D`).
|
||||
"""
|
||||
stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
|
||||
# Tensor Validation
|
||||
BSD = x.shape
|
||||
validate_3d(x, *BSD)
|
||||
validate_weight_bias(weight, BSD[2])
|
||||
validate_weight_bias(bias, BSD[2])
|
||||
validate_3d(scale, *BSD)
|
||||
validate_3d(shift, *BSD)
|
||||
validate_weight_bias(weight2, BSD[2])
|
||||
validate_weight_bias(bias2, BSD[2])
|
||||
validate_3d(scale2, *BSD)
|
||||
if norm_type == "layer" or norm_type == "rms":
|
||||
D = x.shape[-1]
|
||||
if D % 256 != 0 or D > 8192:
|
||||
raise ValueError(
|
||||
f"D={D} not supported, must be multiple of 256 and <= 8192"
|
||||
)
|
||||
y = torch.empty_like(x) # create output tensor
|
||||
y2 = torch.empty_like(x) # create output tensor
|
||||
scale = broadcast_tensor_for_bsfd(scale, *x.shape) # handle various shapes
|
||||
shift = broadcast_tensor_for_bsfd(shift, *x.shape) # handle various shapes
|
||||
scale2 = broadcast_tensor_for_bsfd(scale2, *x.shape) # handle various shapes
|
||||
torch_tensors = [y, y2, x, weight, bias, scale, shift, weight2, bias2, scale2]
|
||||
cute_tensor_args = [to_cute_arg(t) for t in torch_tensors]
|
||||
# Compile cache
|
||||
hash_key = NormTanhMulAddNormScale.make_hash_key(norm_type, *torch_tensors)
|
||||
compiled_fn = _COMPILE_CACHE.get(hash_key)
|
||||
if compiled_fn is None:
|
||||
kernel = NormTanhMulAddNormScale(D, norm_type, is_norm2=True)
|
||||
fake_sig_args = [to_fake_cute_args(t) for t in torch_tensors]
|
||||
compiled_fn = cute.compile(
|
||||
kernel, *fake_sig_args, options="--enable-tvm-ffi"
|
||||
)
|
||||
_COMPILE_CACHE[hash_key] = compiled_fn
|
||||
# Execute
|
||||
compiled_fn(*cute_tensor_args, eps, stream)
|
||||
return y, y2
|
||||
else:
|
||||
raise ValueError(f'norm_type must be one of "layer" and "rms"')
|
||||
|
||||
|
||||
@fused_norm_tanh_mul_add_norm_scale.register_fake
|
||||
def _fused_norm_tanh_mul_add_norm_scale_fake(
|
||||
x, weight, bias, scale, shift, weight2, bias2, scale2, norm_type, eps=1e-5
|
||||
):
|
||||
return x.new_empty(x.shape), x.new_empty(x.shape)
|
||||
@@ -0,0 +1,88 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
|
||||
fused_norm_tanh_mul_add,
|
||||
fused_norm_tanh_mul_add_norm_scale,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=45, suite="stage-b-kernel-unit-1-gpu-large")
|
||||
register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
BSD_CONFIG = [
|
||||
(1, 3648, 3840), # Z-image
|
||||
(1, 4128, 3840), # Z-image
|
||||
(3, 7, 256), # bound
|
||||
(7, 1, 8192), # bound
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B,S,D", BSD_CONFIG)
|
||||
@pytest.mark.parametrize("norm_type", ["rms", "layer"])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_norm_tanh_mul_add(B: int, S: int, D: int, norm_type: str, dtype: str) -> None:
|
||||
device = "cuda"
|
||||
eps = 1e-5
|
||||
x = torch.randn(B, S, D, device=device, dtype=dtype)
|
||||
weight = torch.randn(D, device=device, dtype=dtype)
|
||||
bias = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None
|
||||
scale = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
shift = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
|
||||
y = fused_norm_tanh_mul_add(x, weight, bias, scale, shift, norm_type, eps)
|
||||
if norm_type == "rms":
|
||||
normed = torch.rms_norm(x, x.shape[-1:], weight=weight, eps=eps)
|
||||
else:
|
||||
normed = torch.layer_norm(x, x.shape[-1:], weight=weight, bias=bias, eps=eps)
|
||||
ref_y = normed * torch.tanh(scale) + shift
|
||||
# Accuracy check
|
||||
if dtype == "float32":
|
||||
torch.testing.assert_close(y, ref_y, atol=1e-5, rtol=1e-5)
|
||||
else:
|
||||
torch.testing.assert_close(y, ref_y, atol=5e-2, rtol=5e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("B,S,D", BSD_CONFIG)
|
||||
@pytest.mark.parametrize("norm_type", ["rms", "layer"])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_norm_tanh_mul_add_norm_scale(
|
||||
B: int, S: int, D: int, norm_type: str, dtype: str
|
||||
) -> None:
|
||||
device = "cuda"
|
||||
eps = 1e-5
|
||||
x = torch.randn(B, S, D, device=device, dtype=dtype)
|
||||
weight = torch.randn(D, device=device, dtype=dtype)
|
||||
bias = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None
|
||||
scale = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
shift = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
weight2 = torch.randn(D, device=device, dtype=dtype)
|
||||
bias2 = torch.randn(D, device=device, dtype=dtype) if norm_type == "layer" else None
|
||||
scale2 = torch.randn(B, 1, D, device=device, dtype=dtype)
|
||||
|
||||
y, y2 = fused_norm_tanh_mul_add_norm_scale(
|
||||
x, weight, bias, scale, shift, weight2, bias2, scale2, norm_type, eps
|
||||
)
|
||||
if norm_type == "rms":
|
||||
normed = torch.rms_norm(x, x.shape[-1:], weight=weight, eps=eps)
|
||||
else:
|
||||
normed = torch.layer_norm(x, x.shape[-1:], weight=weight, bias=bias, eps=eps)
|
||||
ref_y = normed * torch.tanh(scale) + shift
|
||||
if norm_type == "rms":
|
||||
normed2 = torch.rms_norm(ref_y, ref_y.shape[-1:], weight=weight2, eps=eps)
|
||||
else:
|
||||
normed2 = torch.layer_norm(
|
||||
ref_y, ref_y.shape[-1:], weight=weight2, bias=bias2, eps=eps
|
||||
)
|
||||
ref_y2 = normed2 * (1 + scale2)
|
||||
# Accuracy check
|
||||
if dtype == "float32":
|
||||
torch.testing.assert_close(y, ref_y, atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(y2, ref_y2, atol=1e-5, rtol=1e-5)
|
||||
else:
|
||||
torch.testing.assert_close(y, ref_y, atol=5e-2, rtol=5e-2)
|
||||
torch.testing.assert_close(y2, ref_y2, atol=5e-2, rtol=5e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
@@ -530,6 +530,80 @@ class RMSNormScaleShift(_NormScaleShift):
|
||||
norm_type = "rms"
|
||||
|
||||
|
||||
################################################################################
|
||||
# NormTanhMulAdd
|
||||
# y = norm(x) * tanh(scale) + shift (where norm is layernorm or rmsnorm)
|
||||
# See details in norm_tanh_mul_add_norm_scale.py
|
||||
################################################################################
|
||||
class _NormTanhMulAdd(CustomOp):
|
||||
norm_type: str
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
eps: float = 1e-6,
|
||||
affine: bool = False,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
):
|
||||
super().__init__()
|
||||
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=affine, eps=eps, dtype=dtype
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(f"Norm type {self.norm_type} not implemented")
|
||||
|
||||
def forward_cuda(
|
||||
self, x: torch.Tensor, scale: torch.Tensor, shift: 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, scale, shift)
|
||||
|
||||
from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
|
||||
fused_norm_tanh_mul_add,
|
||||
)
|
||||
|
||||
x, scale, shift = x.contiguous(), scale.contiguous(), shift.contiguous()
|
||||
weight = _ensure_contiguous(getattr(self.norm, "weight", None))
|
||||
bias = _ensure_contiguous(getattr(self.norm, "bias", None))
|
||||
return fused_norm_tanh_mul_add(
|
||||
x,
|
||||
weight,
|
||||
bias,
|
||||
scale,
|
||||
shift,
|
||||
self.norm_type,
|
||||
self.eps,
|
||||
)
|
||||
|
||||
def forward_hip(self, *args, **kwargs):
|
||||
# Fallback to native because ROCm does not support CuTeDSL.
|
||||
return self.forward_native(*args, **kwargs)
|
||||
|
||||
def forward_native(
|
||||
self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
y = self.norm(x) * torch.tanh(scale) + shift
|
||||
return y.to(x.dtype)
|
||||
|
||||
|
||||
class LayerNormTanhMulAdd(_NormTanhMulAdd):
|
||||
norm_type = "layer"
|
||||
|
||||
|
||||
class RMSNormTanhMulAdd(_NormTanhMulAdd):
|
||||
norm_type = "rms"
|
||||
|
||||
|
||||
def apply_qk_norm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
@@ -709,6 +783,34 @@ def apply_qk_norm_rope(
|
||||
)
|
||||
|
||||
|
||||
def apply_rmsnorm_tanh_mul_add(
|
||||
x: torch.Tensor,
|
||||
gate: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
norm: "RMSNorm",
|
||||
) -> torch.Tensor:
|
||||
"""Compute residual + tanh(gate) * rmsnorm(x), with a fused CUDA fast path."""
|
||||
if get_bool_env_var("SGLANG_ENABLE_DETERMINISTIC_INFERENCE"):
|
||||
return residual + torch.tanh(gate) * norm(x)
|
||||
|
||||
if _is_cuda and x.is_cuda and x.shape[-1] % 256 == 0 and x.shape[-1] <= 8192:
|
||||
from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
|
||||
fused_norm_tanh_mul_add,
|
||||
)
|
||||
|
||||
return fused_norm_tanh_mul_add(
|
||||
x.contiguous(),
|
||||
norm.weight.data.contiguous(),
|
||||
None,
|
||||
gate.contiguous(),
|
||||
residual.contiguous(),
|
||||
"rms",
|
||||
norm.variance_epsilon,
|
||||
)
|
||||
|
||||
return residual + torch.tanh(gate) * norm(x)
|
||||
|
||||
|
||||
def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor:
|
||||
tp_rank = get_tensor_model_parallel_rank()
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
|
||||
@@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.layers.attention import (
|
||||
from sglang.multimodal_gen.runtime.layers.layernorm import (
|
||||
RMSNorm,
|
||||
apply_qk_norm_with_optional_rope,
|
||||
apply_rmsnorm_tanh_mul_add,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
@@ -435,8 +436,7 @@ class ZImageTransformerBlock(nn.Module):
|
||||
scale_msa, gate_msa, scale_mlp, gate_mlp = scale_msa_gate.unsqueeze(
|
||||
1
|
||||
).chunk(4, dim=2)
|
||||
gate_msa, gate_mlp = gate_msa.tanh(), gate_mlp.tanh()
|
||||
scale_msa, scale_mlp = 1.0 + scale_msa, 1.0 + scale_mlp
|
||||
scale_msa = 1.0 + scale_msa
|
||||
|
||||
# Attention block
|
||||
attn_out = self.attention(
|
||||
@@ -445,14 +445,39 @@ class ZImageTransformerBlock(nn.Module):
|
||||
num_replicated_prefix=num_replicated_prefix,
|
||||
num_replicated_suffix=num_replicated_suffix,
|
||||
)
|
||||
x = x + gate_msa * self.attention_norm2(attn_out)
|
||||
if (
|
||||
_is_cuda
|
||||
and attn_out.is_cuda
|
||||
and attn_out.shape[-1] % 256 == 0
|
||||
and attn_out.shape[-1] <= 8192
|
||||
and self.attention_norm2.variance_epsilon
|
||||
== self.ffn_norm1.variance_epsilon
|
||||
):
|
||||
from sglang.jit_kernel.diffusion.cutedsl.norm_tanh_mul_add_norm_scale import (
|
||||
fused_norm_tanh_mul_add_norm_scale,
|
||||
)
|
||||
|
||||
x, ffn_in = fused_norm_tanh_mul_add_norm_scale(
|
||||
attn_out.contiguous(),
|
||||
self.attention_norm2.weight.data.contiguous(),
|
||||
None,
|
||||
gate_msa.contiguous(),
|
||||
x.contiguous(),
|
||||
self.ffn_norm1.weight.data.contiguous(),
|
||||
None,
|
||||
scale_mlp.contiguous(),
|
||||
"rms",
|
||||
self.attention_norm2.variance_epsilon,
|
||||
)
|
||||
else:
|
||||
x = apply_rmsnorm_tanh_mul_add(
|
||||
attn_out, gate_msa, x, self.attention_norm2
|
||||
)
|
||||
ffn_in = self.ffn_norm1(x) * (1.0 + scale_mlp)
|
||||
|
||||
# FFN block
|
||||
x = x + gate_mlp * self.ffn_norm2(
|
||||
self.feed_forward(
|
||||
self.ffn_norm1(x) * scale_mlp,
|
||||
)
|
||||
)
|
||||
ffn_out = self.feed_forward(ffn_in)
|
||||
x = apply_rmsnorm_tanh_mul_add(ffn_out, gate_mlp, x, self.ffn_norm2)
|
||||
else:
|
||||
# Attention block
|
||||
attn_input = self.attention_norm1(x)
|
||||
|
||||
Reference in New Issue
Block a user