From 9b23589c80612d16d3a37eb47910dbf01045587e Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Tue, 7 Apr 2026 10:57:02 +0800 Subject: [PATCH] [cherry-pick][Perf] Skip Wan2.2 cross attn Ulysses SP (#2459) (#2537) Signed-off-by: gcanlin Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Canlin Guo --- 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(