[v0.18.0.post1][Perf] Add layout field of attention metadata for short shape of attn

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

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)

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -187,7 +187,16 @@ 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
+        layout = attn_metadata.layout if attn_metadata else None
+
         try:
             from mindiesd import attention_forward
         except ImportError:
@@ -198,14 +207,13 @@ 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,
             value,
             attn_mask=attention_mask,
             opt_mode="manual",
             op_type="fused_attn_score",
-            layout="BNSD",
+            layout=layout,
         )
         return output

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)