[Diffusion][NPU] Add support for MOVA (#21633)

Co-authored-by: zhangshuai (S) <z00836796@china.huawei.com>

python/sglang/jit_kernel/diffusion/triton/npu_fallback.py

@@ -1,4 +1,8 @@
 import torch
+import torch_npu
+
+NPU_ROTARY_MUL_MAX_NUM_HEADS = 1000
+NPU_ROTARY_MUL_MAX_HEAD_SIZE = 896
 
 
 # TODO: remove this when triton ascend bug is fixed
@@ -18,6 +22,23 @@ def apply_rotary_embedding_native(
 ) -> torch.Tensor:
     cos = cos.unsqueeze(-2).to(x.dtype)
     sin = sin.unsqueeze(-2).to(x.dtype)
+
+    if (
+        cos.dim() == 3
+        and x.dim() == 3
+        and x.shape[1] < NPU_ROTARY_MUL_MAX_NUM_HEADS
+        and x.shape[2] < NPU_ROTARY_MUL_MAX_HEAD_SIZE
+    ):
+        if cos.size(-1) * 2 == x.size(-1):
+            cos = torch.cat([cos, cos], dim=-1)
+            sin = torch.cat([sin, sin], dim=-1)
+        cos = cos.unsqueeze(0)
+        sin = sin.unsqueeze(0)
+        x = x.unsqueeze(0)
+        x_embed = torch_npu.npu_rotary_mul(x, cos, sin)
+        x_embed = x_embed.squeeze(0)
+        return x_embed
+
     x1 = x[..., ::2]
     x2 = x[..., 1::2]
     o1 = x1 * cos - x2 * sin