[Hardware][Ascend] Integrate MindIE dense and sparse quantized attention for Wan2.2 T2V (#7561)

Signed-off-by: zqxu <xuzhenqiang3@huawei.com>
Co-authored-by: zqxu <xuzhenqiang3@huawei.com>
This commit is contained in:
Zhenqiang Xuandzqxu authored and GitHub committed 2026-09-18 01:28:27 +00:00
1 parent f8197e120f
commit 242c3cddeb
15 files changed
+853 -368

No files matched your search

@@ -0,0 +1,40 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Wan integration contracts required by quantized dense/sparse attention."""
from unittest.mock import Mock
import pytest
import torch
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformerBlock
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_block_publishes_unpadded_grid_without_partial_packed_metadata():
block = object.__new__(WanTransformerBlock)
torch.nn.Module.__init__(block)
block.scale_shift_table = torch.zeros(1, 6, 4)
block.norm1 = lambda x, *a: x
block.norm2 = lambda x: x
block.norm3 = lambda x, *a: x
block.attn1 = Mock(side_effect=lambda x, *a: torch.zeros_like(x))
block.attn2 = Mock(side_effect=lambda x, *a: torch.zeros_like(x))
block.ffn = lambda x: torch.zeros_like(x)
# Local SP tokens include alignment padding; grid describes the full valid video.
x = torch.zeros(1, 1152, 4)
block(
x,
x,
torch.zeros(1, 6, 4),
None,
hidden_states_mask=torch.ones(1, 1152, dtype=torch.bool),
vsa_dit_seq_shape=(5, 16, 27),
)
metadata = block.attn1.call_args.args[2]
assert metadata.video_layout.latent_grid == (5, 16, 27)
assert metadata.video_layout.used_len == 2160
assert metadata.video_layout.prefix_len == 0
assert "max_seqlen_q" not in metadata.extra # CUDA must not infer incomplete packed varlen.
assert block.attn2.call_args.args[2] is None