From 2f22ed58ea907802abbaf6145a9236f6c86f9e7f Mon Sep 17 00:00:00 2001
From: Seraphim Volochaev <116020688+Svoloch2940194@users.noreply.github.com>
Date: Wed, 5 Aug 2026 04:06:00 -0700
Subject: [PATCH] [NPU] Adding a fast layernorm for diffusion models and fix
 BSA (#29027)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
---
 .../attention/backends/block_sparse_attn.py   |  2 +-
 .../runtime/layers/layernorm.py               | 58 ++++++++++++++++++-
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py
index 75f56b73e..fa6df1ef2 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/block_sparse_attn.py
@@ -187,7 +187,7 @@ class BlockSparseAttentionImpl(AttentionImpl):
         smask: torch.Tensor,
         sct: torch.Tensor,
     ) -> torch.Tensor:
-        return torch.ops.attentions.block_sparse_attention(
+        return torch.ops.attentions.ada_block_sparse_attention(
             query=query,
             key=key,
             value=value,
diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
index 87134a55b..fe0eb30d3 100755
--- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py
+++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
@@ -425,7 +425,42 @@ class LayerNorm(CustomOp):
 # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
 # NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using
 # FSDP's MixedPrecisionPolicy
-class FP32LayerNorm(nn.LayerNorm):
+@CustomOp.register("fp32_layer_norm")
+class FP32LayerNorm(CustomOp, nn.LayerNorm):
+
+    def __init__(
+        self,
+        normalized_shape,
+        eps=1e-5,
+        elementwise_affine=True,
+        bias=True,
+        device=None,
+        dtype=None,
+    ):
+        nn.LayerNorm.__init__(
+            self,
+            normalized_shape=normalized_shape,
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+            bias=bias,
+            device=device,
+            dtype=dtype,
+        )
+        self._forward_method = self.dispatch_forward()
+
+        try:
+            import attentions  # noqa: F401
+        except ImportError:
+            from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+            logger = init_logger(__name__)  # pylint: disable=invalid-name
+            logger.warning(
+                "The 'attentions' library is not installed. Falling back to native layernorm. "
+                "Installing this library may improve performance on NPU."
+                "See: sgl-project/sgl-kernel-npu"
+            )
+            self._forward_method = self.forward_native
+
     def _cached_fp32_param(
         self, attr: str, param: torch.Tensor | None, device: torch.device
     ) -> torch.Tensor | None:
@@ -452,7 +487,7 @@ class FP32LayerNorm(nn.LayerNorm):
         self.__dict__[attr] = (key, fp32_param)
         return fp32_param
 
-    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+    def forward_native(self, inputs: torch.Tensor) -> torch.Tensor:
         origin_dtype = inputs.dtype
         device = inputs.device
         weight = self._cached_fp32_param("_weight_fp32_cache", self.weight, device)
@@ -465,6 +500,25 @@ class FP32LayerNorm(nn.LayerNorm):
             self.eps,
         ).to(origin_dtype)
 
+    def forward_cuda(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(inputs)
+
+    def forward_npu(self, inputs: torch.Tensor) -> torch.Tensor:
+        origin_dtype = inputs.dtype
+        device = inputs.device
+        weight = self._cached_fp32_param("_weight_fp32_cache", self.weight, device)
+        bias = self._cached_fp32_param("_bias_fp32_cache", self.bias, device)
+
+        output, _, _ = torch.ops.attentions.layernorm(
+            input=inputs,
+            normalized_shape=list(self.normalized_shape),
+            weight=weight,
+            bias=bias,
+            eps=self.eps,
+            impl_mode=0,
+        )
+        return output.to(origin_dtype)
+
 
 ################################################################################
 # Fused norm kernel