[Diffusion][NPU] Add support for MOVA

python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py

@@ -7,8 +7,10 @@
 from sglang.jit_kernel.diffusion.triton.rotary import apply_rotary_embedding
 from sglang.kernel_api_logging import debug_kernel_api
 from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.srt.layers.rotary_embedding.utils import apply_rotary_pos_emb
 from sglang.srt.utils.custom_op import register_custom_op_from_extern
 
+_is_npu = current_platform.is_npu()
 _is_cuda = current_platform.is_cuda()
 if _is_cuda:
     try:
@@ -111,8 +113,15 @@ def apply_flashinfer_rope_qk_inplace(
             sin = cos_sin_cache[positions, half_size:].to(q.dtype)
         q_flat = q.reshape(bsz * seqlen, nheads, d)
         k_flat = k.reshape(bsz * seqlen, nheads, d)
-        q_rot = apply_rotary_embedding(q_flat, cos, sin, interleaved=not is_neox)
-        k_rot = apply_rotary_embedding(k_flat, cos, sin, interleaved=not is_neox)
+
+        if _is_npu and is_neox:
+            cos = torch.cat([cos, cos], dim=-1)
+            sin = torch.cat([sin, sin], dim=-1)
+            q_rot, k_rot = apply_rotary_pos_emb(q_flat, k_flat, cos, sin)
+        else:
+            q_rot = apply_rotary_embedding(q_flat, cos, sin, interleaved=not is_neox)
+            k_rot = apply_rotary_embedding(k_flat, cos, sin, interleaved=not is_neox)
+
         return q_rot.view(bsz, seqlen, nheads, d), k_rot.view(bsz, seqlen, nheads, d)
 
     if positions is None:

python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py

@@ -16,6 +16,7 @@
 from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
 from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
 from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention
+from sglang.multimodal_gen.runtime.platforms import current_platform
 
 # Reuse SGLang's optimized RMSNorm instead of torch.nn.RMSNorm or custom SlowRMSNorm
 from sglang.multimodal_gen.runtime.layers.layernorm import (
@@ -522,7 +523,10 @@ def patchify(
         self, x: torch.Tensor, control_camera_latents_input: torch.Tensor | None = None
     ):
         # NOTE(dhyu): avoid slow_conv
-        x = x.contiguous(memory_format=torch.channels_last_3d)
+        if current_platform.is_npu:
+            x = x.contiguous()
+        else:
+            x = x.contiguous(memory_format=torch.channels_last_3d)
         x = self.patch_embedding(x)
         grid_size = x.shape[2:]
         x = rearrange(x, "b c f h w -> b (f h w) c").contiguous()

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py

@@ -69,6 +69,7 @@
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 from sglang.srt.utils.common import get_compiler_backend
 
+_is_npu = current_platform.is_npu()
 logger = init_logger(__name__)
 
 
@@ -714,7 +715,12 @@ def inference_single_step(
 
         # Build visual freqs for full sequence
         visual_dit._init_freqs()
-        visual_freqs = tuple(freq.to(visual_x.device) for freq in visual_dit.freqs)
+        if _is_npu:
+            visual_freqs = tuple(
+                freq.to(device=visual_x.device, dtype=torch.complex64)
+                for freq in visual_dit.freqs)
+        else:
+            visual_freqs = tuple(freq.to(visual_x.device) for freq in visual_dit.freqs)
         visual_freqs = (
             torch.cat(
                 [
@@ -734,17 +740,22 @@ def inference_single_step(
 
         # Build audio freqs for full sequence
         self.audio_dit._init_freqs()
+        if _is_npu:    
+            audio_freqs = tuple(
+                freq.to(device=audio_x.device, dtype=torch.complex64)
+                for freq in self.audio_dit.freqs)
+        else:
+            audio_freqs = tuple(freq.to(audio_x.device) for freq in self.audio_dit.freqs)
         audio_freqs = (
             torch.cat(
                 [
-                    self.audio_dit.freqs[0][:f].view(f, -1).expand(f, -1),
-                    self.audio_dit.freqs[1][:f].view(f, -1).expand(f, -1),
-                    self.audio_dit.freqs[2][:f].view(f, -1).expand(f, -1),
+                    audio_freqs[0][:f].view(f, -1).expand(f, -1),
+                    audio_freqs[1][:f].view(f, -1).expand(f, -1),
+                    audio_freqs[2][:f].view(f, -1).expand(f, -1),
                 ],
                 dim=-1,
             )
             .reshape(full_audio_seq_len, 1, -1)
-            .to(audio_x.device)
         )
 
         # Shard sequences for SP