[gcanlin]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

ascend_laser_attention (LA) doesn't support different Q/KV sequence lengths, which causes errors for cross attention where Q comes from decoder (long seq) and KV comes from encoder (short seq, e.g. 512 text tokens).

This change:

• Detects Q/KV seqlen mismatch in forward_npu and bypasses mindiesd to use torch_npu.npu_fusion_attention directly with fused_attn_score
• Auto-detects input layout (BSND/BNSD) based on num_heads position
• Self attention (same Q/KV seqlen) still respects MINDIE_SD_FA_TYPE env var
• Reverts WanCrossAttention to use unified Attention layer

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[gcanlin]

lint format.

[SamitHuang]

can you report the performance measurements in the description? e.g. run benchmarks comparing the original Attention layer with this SDPAImpl workaround on realistic workloads and report the generation latency and peak VRAM usage before and after. also verify that quality is preserved without accuracy degradation.

[SamitHuang]

Bypassing the unified Attention layer skips the pre_attention and post_attention hooks used for parallel strategies. This breaks the abstraction and may cause issues with distributed execution. Consider adding a backend_override parameter to the Attention class instead of hardcoding SDPAImpl here.

[gcanlin]

We expect to only change Wan2.2 modeling file. If add backend_override in attention, we have to change the attention layer. As a workaround, I prefer to hardcode here. We have confirmed the performance didn't happen degradation offline.

[hsliuustc0106]

BLOCKER scan:

• Correctness: PASS
• Reliability/Safety: PASS
• Breaking Changes: PASS
• Test Coverage: PASS (existing Wan2.2 e2e tests cover this)
• Documentation: PASS (workaround, no docs needed)
• Security: PASS

OVERALL: NO BLOCKERS

VERDICT: COMMENT

Question: This is marked as a [Workaround]. Is there a plan to add SDPA as a configurable option with proper fallback to other backends for different hardware? The current forced SDPA approach may not work on all GPUs.

[gcanlin]

have updated to only be applied on NPU.

[lishunyang12]

Review of [v0.18.0.post1][Perf] handle cross attention with different Q/KV seqlen on NPU

Overall the approach is sound: bypass LA (which does not support mismatched Q/KV seqlen) and call torch_npu.npu_fusion_attention directly for cross attention. A few items worth addressing:

1. Layout detection heuristic is fragile (main concern)

if query.shape[2] == self.num_heads:
    layout = "BSND"
else:
    layout = "BNSD"

This assumes S != num_heads. If a model happens to use S == num_heads (e.g., 24 heads with a 24-token sequence), the layout will be misidentified. Consider one of:

• Accept layout as a constructor / method parameter so callers declare it explicitly.
• Add a secondary check, e.g., also verify query.shape[1] != self.num_heads and raise an AssertionError if both dimensions match, making the ambiguity visible.
• At minimum, add a comment documenting this assumption and the known limitation.

2. pre_tokens / next_tokens magic numbers

pre_tokens=2147483647,
next_tokens=2147483647,

These are INT32_MAX. They effectively disable causal masking, which is correct for cross attention, but they should be named constants (e.g., _NPU_NO_CAUSAL = 2147483647) with a one-line comment explaining why, so future readers do not wonder if this is an overflow bug.

3. import torch_npu inside the hot path

The import torch_npu is inside the if q_seqlen != k_seqlen branch, which runs on every cross-attention forward call. Python caches module imports after the first load so the runtime cost is small, but it is conventional to do the import once at module level (guarded by a try/except) or at __init__ time, similar to how mindiesd is imported in the self-attention path. This also lets you fail fast at startup if torch_npu is missing.

4. PR description vs. diff mismatch

The description states "Reverts WanCrossAttention to use unified Attention layer", but no changes to wan2_2_transformer.py (or any Wan model file) are included in this diff. Was that change dropped or is it in a separate PR?

5. Test coverage

The test file change is just an import style fix (combining two from fastapi import ... lines). There are no new tests exercising the cross-attention path. Given this is a correctness-sensitive code path, at least a unit test verifying that:

• BSND and BNSD layouts are detected correctly,
• the torch_npu path is taken when q_seqlen != k_seqlen,
• the mindiesd path is taken when q_seqlen == k_seqlen,

would be valuable, even if they need to be mocked for CI environments without NPU hardware.

Minor

• The attention_forward call in the self-attention path now uses the dynamically detected layout instead of the hardcoded "BNSD". This is a good improvement, but please confirm that attention_forward from mindiesd supports "BSND" layout -- the old code always passed "BNSD".

---

Not blocking, but items 1 and 4 should be addressed before merge.

[gcanlin]

@SamitHuang @david6666666 PTAL. A new workaround. Will integrate it into #2681 using an elegant way.