From 95a07f7732900974d9e608b39f36e5b2e6518442 Mon Sep 17 00:00:00 2001 From: "Hui." <62495465+Hu1Lcode@users.noreply.github.com> Date: Thu, 23 Apr 2026 17:47:43 +0800 Subject: [PATCH] [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 | 60 +++++++++++++++++-- .../models/wan2_2/wan2_2_transformer.py | 48 +++++---------- .../models/wan2_2/wan2_2_vace_transformer.py | 13 +++- 3 files changed, 83 insertions(+), 38 deletions(-) diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 61ddb4d84..127e1b2cb 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/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, diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index d4d81b78e..98e10b408 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/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) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index c48938e1b..5060f1904 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/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)