[PERF]use mindiesd fused rope and rope cache (#2571)

Signed-off-by: Hui <1779066624@qq.com>
Signed-off-by: Hui. <62495465+Hu1Lcode@users.noreply.github.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>

vllm_omni/diffusion/layers/rope.py

@@ -72,18 +72,18 @@ class RotaryEmbedding(CustomOp):
            of 1st half and 2nd half (GPT-NeoX style).
     """
 
-    def __init__(
-        self,
-        is_neox_style: bool = False,
-    ) -> None:
+    def __init__(self, is_neox_style: bool = False) -> None:
         super().__init__()
         self.is_neox_style = is_neox_style
         self.interleaved = not is_neox_style
         self.apply_rotary_emb_flash_attn = None
+        self.has_mindie = False
         if find_spec("flash_attn") is not None:
             from flash_attn.ops.triton.rotary import apply_rotary
 
             self.apply_rotary_emb_flash_attn = apply_rotary
+        if find_spec("mindiesd") is not None:
+            self.has_mindie = True
 
     def forward_cuda(
         self,
@@ -132,7 +132,7 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
-        if find_spec("mindiesd"):
+        if self.has_mindie:
             return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved)
         else:
             return self.forward_native(x, cos, sin)
@@ -167,6 +167,56 @@ class RotaryEmbedding(CustomOp):
         )
 
 
+class RotaryEmbeddingWan(RotaryEmbedding):
+    """
+    rotary positional embedding for Wan.
+    interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+           of 1st half and 2nd half (GPT-NeoX style).
+    """
+
+    def __init__(self, is_neox_style: bool = False, half_head_dim: bool = False) -> None:
+        super().__init__(is_neox_style=is_neox_style)
+        self.half_head_dim = half_head_dim
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x, cos, sin)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.has_mindie:
+            if cos.dim() > 2:
+                cos = cos.reshape(-1, cos.shape[-1])
+                sin = sin.reshape(-1, sin.shape[-1])
+            return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved, self.half_head_dim)
+        else:
+            return self.forward_native(x, cos, sin)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        x1, x2 = x.unflatten(-1, (-1, 2)).unbind(-1)
+        rotated = torch.stack(
+            (
+                x1 * cos - x2 * sin,
+                x1 * sin + x2 * cos,
+            ),
+            dim=-1,
+        )
+        return rotated.flatten(-2, -1).to(x.dtype)
+
+
 def apply_rope_to_qk(
     rope: RotaryEmbedding,
     query: torch.Tensor,

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -30,40 +30,12 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
 from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm
+from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
 
-def apply_rotary_emb_wan(
-    hidden_states: torch.Tensor,
-    freqs_cos: torch.Tensor,
-    freqs_sin: torch.Tensor,
-) -> torch.Tensor:
-    """
-    Apply rotary embeddings to input tensors using the given frequency tensors.
-
-    Args:
-        hidden_states: Input tensor of shape [B, S, H, D]
-        freqs_cos: Cosine frequencies
-        freqs_sin: Sine frequencies
-
-    Returns:
-        Tensor with rotary embeddings applied
-    """
-    x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
-    cos = freqs_cos[..., 0::2]
-    sin = freqs_sin[..., 1::2]
-    rotated = torch.stack(
-        (
-            x1 * cos - x2 * sin,
-            x1 * sin + x2 * cos,
-        ),
-        dim=-1,
-    )
-    return rotated.flatten(-2, -1).to(hidden_states.dtype)
-
-
 class DistributedRMSNorm(nn.Module):
     """
     RMSNorm that computes global RMS across tensor parallel ranks.
@@ -427,9 +399,10 @@ class WanSelfAttention(nn.Module):
 
         # Apply rotary embeddings
         if rotary_emb is not None:
+            self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
             freqs_cos, freqs_sin = rotary_emb
-            query = apply_rotary_emb_wan(query, freqs_cos, freqs_sin)
-            key = apply_rotary_emb_wan(key, freqs_cos, freqs_sin)
+            query = self.rotary_embedding(query, freqs_cos, freqs_sin)
+            key = self.rotary_embedding(key, freqs_cos, freqs_sin)
 
         # Create attention metadata if mask is provided
         attn_metadata = None
@@ -871,6 +844,10 @@ class WanTransformer3DModel(nn.Module):
         self.timestep_proj_prepare = TimestepProjPrepare()
         self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
 
+        # ROPE helper
+        self._cached_rope_emb = None
+        self._cached_rope_resolution = None
+
     @property
     def dtype(self) -> torch.dtype:
         """Return the dtype of the model parameters."""
@@ -892,7 +869,14 @@ class WanTransformer3DModel(nn.Module):
         post_patch_width = width // p_w
 
         # Compute RoPE embeddings (sharded by _sp_plan via split_output=True)
-        rotary_emb = self.rope(hidden_states)
+        current_rope_resolution = (post_patch_num_frames, post_patch_height, post_patch_width)
+        if self._cached_rope_resolution == current_rope_resolution and self._cached_rope_emb is not None:
+            rotary_emb = self._cached_rope_emb
+        else:
+            freqs_cos, freqs_sin = self.rope(hidden_states)
+            rotary_emb = (freqs_cos[..., 0::2].to(hidden_states.dtype), freqs_sin[..., 1::2].to(hidden_states.dtype))
+            self._hidden_states_shape = hidden_states.shape
+            self._cached_rope_emb = rotary_emb
 
         # Patch embedding and flatten to sequence
         # (hidden_states is sharded at blocks.0 input by _sp_plan)

vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py

@@ -123,6 +123,10 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
                 ]
             )
 
+        # ROPE helper
+        self._cached_rope_emb = None
+        self._cached_rope_resolution = None
+
     def embed_vace_context(
         self,
         vace_context: torch.Tensor,
@@ -165,7 +169,14 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
         post_patch_width = width // p_w
 
         # Compute RoPE embeddings (sharded by _sp_plan via split_output=True)
-        rotary_emb = self.rope(hidden_states)
+        current_rope_resolution = (post_patch_num_frames, post_patch_height, post_patch_width)
+        if self._cached_rope_resolution == current_rope_resolution and self._cached_rope_emb is not None:
+            rotary_emb = self._cached_rope_emb
+        else:
+            freqs_cos, freqs_sin = self.rope(hidden_states)
+            rotary_emb = (freqs_cos[..., 0::2].to(hidden_states.dtype), freqs_sin[..., 1::2].to(hidden_states.dtype))
+            self._hidden_states_shape = hidden_states.shape
+            self._cached_rope_emb = rotary_emb
 
         # Patch embedding and flatten to sequence
         hidden_states = self.patch_embedding(hidden_states)