diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py b/python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py new file mode 100644 index 000000000..c5ae1c619 --- /dev/null +++ b/python/sglang/jit_kernel/diffusion/cutedsl/norm_tanh_mul_add_norm_scale.py @@ -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() + + 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 + ) + 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) diff --git a/python/sglang/jit_kernel/tests/test_norm_tanh_mul_add_norm_scale.py b/python/sglang/jit_kernel/tests/test_norm_tanh_mul_add_norm_scale.py new file mode 100644 index 000000000..e36736917 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_norm_tanh_mul_add_norm_scale.py @@ -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__]) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 092a0b805..71f6b2ce3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -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() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 8a0ab6a58..ebf187e3e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py @@ -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)