From e2892ef6c1d461b65435339ca3ac83c17d8d3c0f Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Thu, 2 Apr 2026 14:28:44 +0800 Subject: [PATCH] [Perf] Optimize Wan2.2 rotary embedding (#2393) Signed-off-by: gcanlin --- .../diffusion/models/wan2_2/wan2_2_transformer.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) 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 a4ae3118a..20e2b9fea 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -52,10 +52,14 @@ def apply_rotary_emb_wan( x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) cos = freqs_cos[..., 0::2] sin = freqs_sin[..., 1::2] - out = torch.empty_like(hidden_states) - out[..., 0::2] = x1 * cos - x2 * sin - out[..., 1::2] = x1 * sin + x2 * cos - return out.type_as(hidden_states) + 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):