From 36c93fc6fb960df01cd04d4ef5ab08302ed4ac50 Mon Sep 17 00:00:00 2001 From: gjsheu Date: Mon, 11 May 2026 12:17:31 +0800 Subject: [PATCH] [NPU] [Diffusion] Use fused operator to improve Wan model E2E performance. (#24028) Co-authored-by: gengjinsong Co-authored-by: sglang-npu-bot Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: gengjinsong <904939979@qq.com> Co-authored-by: ronnie_zheng --- .../runtime/layers/elementwise.py | 7 +++ .../runtime/layers/layernorm.py | 55 ++++++++++++++++++- scripts/ci/npu/npu_ci_install_dependency.sh | 6 +- 3 files changed, 63 insertions(+), 5 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/elementwise.py b/python/sglang/multimodal_gen/runtime/layers/elementwise.py index c990f7f67..181ed143d 100644 --- a/python/sglang/multimodal_gen/runtime/layers/elementwise.py +++ b/python/sglang/multimodal_gen/runtime/layers/elementwise.py @@ -38,3 +38,10 @@ class MulAdd(CustomOp): self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 ): return self.forward_native(a, b, c, k=k) + + def forward_npu( + self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0 + ): + from sgl_kernel_npu.norm.scale_shift import fused_scale_shift + + return fused_scale_shift(a, b, c, scale_constant=k) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index 883e2a808..d5e16c488 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -520,6 +520,38 @@ class _ScaleResidualNormScaleShift(CustomOp): modulated = fuse_scale_shift_kernel(normalized, scale, shift) return modulated, residual_output + def forward_npu( + self, + residual: torch.Tensor, + x: torch.Tensor, + gate: torch.Tensor | int, + shift: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + from sgl_kernel_npu.norm.scale_shift import fused_scale_shift + + # x.shape: [batch_size, seq_len, inner_dim] + if isinstance(gate, int): + # used by cross-attention, should be 1 + assert gate == 1 + residual_output = residual + x + elif isinstance(gate, torch.Tensor): + if gate.dim() == 4: + # gate.shape: [batch_size, num_frames, 1, inner_dim] + num_frames = gate.shape[1] + frame_seqlen = x.shape[1] // num_frames + residual_output = residual + ( + x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate + ).flatten(1, 2) + else: + # gate.shape: [batch_size, 1, inner_dim] + residual_output = residual + x * gate + else: + raise ValueError(f"Gate type {type(gate)} not supported") + normalized = self.norm(residual_output) + modulated = fused_scale_shift(normalized, scale, shift) + return modulated, residual_output + class ScaleResidualLayerNormScaleShift(_ScaleResidualNormScaleShift): norm_type = "layer" @@ -606,6 +638,15 @@ class _NormScaleShift(CustomOp): modulated = fuse_scale_shift_kernel(normalized, scale, shift) return modulated.to(x.dtype) + def forward_npu( + self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor + ) -> torch.Tensor: + from sgl_kernel_npu.norm.scale_shift import fused_scale_shift + + normalized = self.norm(x) + modulated = fused_scale_shift(normalized, scale, shift) + return modulated.to(x.dtype) + class LayerNormScaleShift(_NormScaleShift): norm_type = "layer" @@ -902,9 +943,19 @@ def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor: src_dtype = x.dtype weight = norm.weight.tensor_split(tp_size)[tp_rank].float() x_fp32 = x.float() - variance = x_fp32.pow(2).mean(dim=-1, keepdim=True) + if _is_npu: + from sgl_kernel_npu.norm.rmsnorm_split import fused_rsqrt_mul, fused_variance + + variance = fused_variance(x_fp32) + else: + variance = x_fp32.pow(2).mean(dim=-1, keepdim=True) + variance = get_tp_group().all_reduce( variance, op=torch._C._distributed_c10d.ReduceOp.AVG ) - output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight + + if _is_npu: + output = fused_rsqrt_mul(x_fp32, variance, weight, norm.variance_epsilon) + else: + output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight return output.to(dtype=src_dtype) diff --git a/scripts/ci/npu/npu_ci_install_dependency.sh b/scripts/ci/npu/npu_ci_install_dependency.sh index e180cf083..9782861f4 100755 --- a/scripts/ci/npu/npu_ci_install_dependency.sh +++ b/scripts/ci/npu/npu_ci_install_dependency.sh @@ -68,10 +68,10 @@ ${UV_PIP_INSTALL} triton-ascend ### Install sgl-kernel-npu -SGLANG_KERNEL_NPU_TAG="2026.03.10.rc1" +SGLANG_KERNEL_NPU_TAG="2026.05.01" mkdir sgl-kernel-npu -(cd sgl-kernel-npu && wget "${GITHUB_PROXY_URL:=""}https://github.com/sgl-project/sgl-kernel-npu/releases/download/${SGLANG_KERNEL_NPU_TAG}/sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch2.8.0-py311-cann8.5.0-${DEVICE_TYPE}-$(arch).zip" \ -&& unzip ./sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch2.8.0-py311-cann8.5.0-${DEVICE_TYPE}-$(arch).zip \ +(cd sgl-kernel-npu && wget "${GITHUB_PROXY_URL:=""}https://github.com/sgl-project/sgl-kernel-npu/releases/download/${SGLANG_KERNEL_NPU_TAG}/sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch${PYTORCH_VERSION}-py311-cann8.5.0-${DEVICE_TYPE}-$(arch).zip" \ +&& unzip ./sgl-kernel-npu-${SGLANG_KERNEL_NPU_TAG}-torch${PYTORCH_VERSION}-py311-cann8.5.0-${DEVICE_TYPE}-$(arch).zip \ && ${UV_PIP_INSTALL} ./deep_ep*.whl ./sgl_kernel_npu*.whl \ && (cd "$(python3 -m pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so))