From 02a6e317d7739ca173e509d89ac295f0ba40dc5f Mon Sep 17 00:00:00 2001 From: gcanlin Date: Tue, 14 Apr 2026 03:58:32 +0000 Subject: [PATCH 1/4] [Perf][Workaround] Force Wan2.2 cross attn to SDPA Signed-off-by: gcanlin --- .../entrypoints/openai_api/test_video_server.py | 3 +-- .../models/wan2_2/wan2_2_transformer.py | 17 ++++++++--------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 266900f707a..465de101648 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -14,8 +14,7 @@ from types import SimpleNamespace import pytest -from fastapi import FastAPI -from fastapi import HTTPException +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from PIL import Image from pytest_mock import MockerFixture 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 c4e3b40cdd5..9f7a2a28c9d 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -23,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, @@ -444,7 +445,7 @@ def forward( class WanCrossAttention(nn.Module): """ - Optimized cross-attention module using vLLM layers. + Optimized cross-attention module using vLLM layers with SDPA. Handles both text cross-attention and optional image cross-attention (I2V). """ @@ -532,14 +533,12 @@ def __init__( ) self.dropout = nn.Dropout(dropout) - # Unified attention layer - self.attn = Attention( + # Force SDPA backend for cross-attention (workaround) + self.sdpa_impl = SDPAImpl( num_heads=self.num_heads, head_size=head_dim, - num_kv_heads=self.num_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, - skip_sequence_parallel=True, ) def forward( @@ -564,7 +563,7 @@ def forward( value = self.to_v(encoder_hidden_states) key = self.norm_k(key) - # Reshape for multi-head attention + # Reshape for multi-head attention: [B, S, H, D] query = query.unflatten(2, (self.num_heads, self.head_dim)) key = key.unflatten(2, (self.num_heads, self.head_dim)) value = value.unflatten(2, (self.num_heads, self.head_dim)) @@ -579,12 +578,12 @@ def forward( key_img = key_img.unflatten(2, (self.num_heads, self.head_dim)) value_img = value_img.unflatten(2, (self.num_heads, self.head_dim)) - hidden_states_img = self.attn(query, key_img, value_img) + hidden_states_img = self.sdpa_impl.forward(query, key_img, value_img) hidden_states_img = hidden_states_img.flatten(2, 3) hidden_states_img = hidden_states_img.type_as(query) - # Main cross-attention using unified attention layer - hidden_states = self.attn(query, key, value) + # Main cross-attention using SDPA (forced) + hidden_states = self.sdpa_impl.forward(query, key, value) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) From 2b329f31920c6c47b8d48b315c98a86838a45e20 Mon Sep 17 00:00:00 2001 From: gcanlin Date: Tue, 14 Apr 2026 04:02:28 +0000 Subject: [PATCH 2/4] remove unnecessary change Signed-off-by: gcanlin --- vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 9f7a2a28c9d..0e61ae5b076 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -563,7 +563,7 @@ def forward( value = self.to_v(encoder_hidden_states) key = self.norm_k(key) - # Reshape for multi-head attention: [B, S, H, D] + # Reshape for multi-head attention query = query.unflatten(2, (self.num_heads, self.head_dim)) key = key.unflatten(2, (self.num_heads, self.head_dim)) value = value.unflatten(2, (self.num_heads, self.head_dim)) From 41487386d642fc1caca54494da9b928521d3844a Mon Sep 17 00:00:00 2001 From: gcanlin Date: Wed, 15 Apr 2026 16:15:14 +0000 Subject: [PATCH 3/4] fix fa Signed-off-by: gcanlin --- .../attention/backends/flash_attn.py | 44 +++++++++++++++++-- .../models/wan2_2/wan2_2_transformer.py | 15 ++++--- 2 files changed, 49 insertions(+), 10 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 5c586c0631e..f524bf18b21 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -187,7 +187,46 @@ def forward_npu( value: torch.Tensor, attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: - """NPU attention implementation using mindiesd.""" + """NPU attention implementation. + + When Q and KV have different seqlen (cross attention), uses torch_npu directly + to bypass environment variable and force fused_attn_score (LA doesn't support + different Q/KV seqlen). Otherwise uses mindiesd attention_forward which respects + MINDIE_SD_FA_TYPE environment variable. + """ + attention_mask = attn_metadata.attn_mask if attn_metadata else None + + # Detect input layout based on num_heads position + # BSND: [B, S, N, D] -> shape[2] == num_heads + # BNSD: [B, N, S, D] -> shape[1] == num_heads + if query.shape[2] == self.num_heads: + layout = "BSND" + q_seqlen, k_seqlen = query.shape[1], key.shape[1] + else: + layout = "BNSD" + q_seqlen, k_seqlen = query.shape[2], key.shape[2] + + # When Q/KV seqlen differ, LA will fail. Use torch_npu directly to bypass env var. + if q_seqlen != k_seqlen: + import torch_npu + + if attention_mask is not None: + attention_mask = ~attention_mask.to(torch.bool) + + output = torch_npu.npu_fusion_attention( + query, + key, + value, + atten_mask=attention_mask, + input_layout=layout, + scale=self.softmax_scale, + pre_tokens=2147483647, + next_tokens=2147483647, + head_num=self.num_heads, + )[0] + return output + + # Self attention: use mindiesd which respects MINDIE_SD_FA_TYPE env var try: from mindiesd import attention_forward except ImportError: @@ -198,7 +237,6 @@ def forward_npu( "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" ) - attention_mask = attn_metadata.attn_mask if attn_metadata else None output = attention_forward( query, key, @@ -206,6 +244,6 @@ def forward_npu( attn_mask=attention_mask, opt_mode="manual", op_type="fused_attn_score", - layout="BNSD", + layout=layout, ) return output 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 0e61ae5b076..c4e3b40cdd5 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -23,7 +23,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata -from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, @@ -445,7 +444,7 @@ def forward( class WanCrossAttention(nn.Module): """ - Optimized cross-attention module using vLLM layers with SDPA. + Optimized cross-attention module using vLLM layers. Handles both text cross-attention and optional image cross-attention (I2V). """ @@ -533,12 +532,14 @@ def __init__( ) self.dropout = nn.Dropout(dropout) - # Force SDPA backend for cross-attention (workaround) - self.sdpa_impl = SDPAImpl( + # Unified attention layer + self.attn = Attention( num_heads=self.num_heads, head_size=head_dim, + num_kv_heads=self.num_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, + skip_sequence_parallel=True, ) def forward( @@ -578,12 +579,12 @@ def forward( key_img = key_img.unflatten(2, (self.num_heads, self.head_dim)) value_img = value_img.unflatten(2, (self.num_heads, self.head_dim)) - hidden_states_img = self.sdpa_impl.forward(query, key_img, value_img) + hidden_states_img = self.attn(query, key_img, value_img) hidden_states_img = hidden_states_img.flatten(2, 3) hidden_states_img = hidden_states_img.type_as(query) - # Main cross-attention using SDPA (forced) - hidden_states = self.sdpa_impl.forward(query, key, value) + # Main cross-attention using unified attention layer + hidden_states = self.attn(query, key, value) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query) From 50c23c9f76363cdb4b80ab2636064a317439f0c1 Mon Sep 17 00:00:00 2001 From: gcanlin Date: Thu, 16 Apr 2026 16:12:11 +0000 Subject: [PATCH 4/4] BSDN Signed-off-by: gcanlin --- .../diffusion/attention/backends/abstract.py | 2 ++ .../attention/backends/flash_attn.py | 32 +------------------ .../models/wan2_2/wan2_2_transformer.py | 4 ++- 3 files changed, 6 insertions(+), 32 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index d0a62bcd9cc..b78645fe691 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -64,6 +64,8 @@ class AttentionMetadata: # a replicated tensor among processes appended to the front or rear of value, depends the joint_strategy joint_strategy: str = "front" # the strategy to joint the query, key, and value, can be "front" or "rear" + layout: str | None = None + # the layout of the input query, key, and value, can be "BSND" or "BNSD" T = TypeVar("T", bound=AttentionMetadata) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index f524bf18b21..9a2df2baf43 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -195,38 +195,8 @@ def forward_npu( MINDIE_SD_FA_TYPE environment variable. """ attention_mask = attn_metadata.attn_mask if attn_metadata else None + layout = attn_metadata.layout if attn_metadata else None - # Detect input layout based on num_heads position - # BSND: [B, S, N, D] -> shape[2] == num_heads - # BNSD: [B, N, S, D] -> shape[1] == num_heads - if query.shape[2] == self.num_heads: - layout = "BSND" - q_seqlen, k_seqlen = query.shape[1], key.shape[1] - else: - layout = "BNSD" - q_seqlen, k_seqlen = query.shape[2], key.shape[2] - - # When Q/KV seqlen differ, LA will fail. Use torch_npu directly to bypass env var. - if q_seqlen != k_seqlen: - import torch_npu - - if attention_mask is not None: - attention_mask = ~attention_mask.to(torch.bool) - - output = torch_npu.npu_fusion_attention( - query, - key, - value, - atten_mask=attention_mask, - input_layout=layout, - scale=self.softmax_scale, - pre_tokens=2147483647, - next_tokens=2147483647, - head_num=self.num_heads, - )[0] - return output - - # Self attention: use mindiesd which respects MINDIE_SD_FA_TYPE env var try: from mindiesd import attention_forward except ImportError: 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 c4e3b40cdd5..eae45d97229 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -583,8 +583,10 @@ def forward( hidden_states_img = hidden_states_img.flatten(2, 3) hidden_states_img = hidden_states_img.type_as(query) + attn_metadata = AttentionMetadata(layout="BSND") + # Main cross-attention using unified attention layer - hidden_states = self.attn(query, key, value) + hidden_states = self.attn(query, key, value, attn_metadata) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.type_as(query)