From 6dc61c9a20a49e86aaf48ad4a03e7c8cb6b29e34 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 3 Apr 2026 16:23:36 +0800 Subject: [PATCH] [Perf] Skip Wan2.2 cross attn Ulysses SP (#2459) Signed-off-by: gcanlin --- vllm_omni/diffusion/attention/layer.py | 4 ++++ vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py | 1 + 2 files changed, 5 insertions(+) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index f83bb294d..4fdf2ff16 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -36,6 +36,7 @@ class Attention(nn.Module): scatter_idx: int = 2, gather_idx: int = 1, use_sync: bool = False, + skip_sequence_parallel: bool = False, ): super().__init__() self.attn_backend = get_attn_backend(-1) @@ -62,6 +63,7 @@ class Attention(nn.Module): self.gather_idx = gather_idx self.use_sync = use_sync self.causal = causal + self.skip_sequence_parallel = skip_sequence_parallel self.use_ring = False self.ring_pg = None @@ -98,6 +100,8 @@ class Attention(nn.Module): (e.g., in noise_refiner/context_refiner before unified_prepare in Z-Image). This avoids unnecessary SP communication for layers not covered by _sp_plan. """ + if self.skip_sequence_parallel: + return self._no_parallel_strategy if is_forward_context_available(): ctx = get_forward_context() if not ctx.sp_active: 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 20e2b9fea..c4e3b40cd 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -539,6 +539,7 @@ class WanCrossAttention(nn.Module): num_kv_heads=self.num_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, + skip_sequence_parallel=True, ) def forward(