[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>

tests/diffusion/attention/test_mindie_public_integration.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cross-repository CPU dispatch contracts, requiring the matching MindIE-SD source.
+
+Only native operators are replaced. Public adapters, rotations, layouts, padding,
+scale preparation, sequence metadata and output cropping execute their real code.
+Operator substitutes carry floating tensors; these are not numerical quantization tests.
+"""
+
+import importlib
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends import flash_attn, rainfusion_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.platforms.npu.quant.kv_quant_npu import get_quant_attention_rotation
+
+mindiesd = pytest.importorskip("mindiesd")
+torch_npu = importlib.import_module("torch_npu")
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "precision,batch,q_len,kv_len,q_heads",
+    [
+        ("fp8", 1, 130, 130, 2),
+        ("mxfp8", 1, 130, 130, 2),
+        ("mxfp4", 1, 130, 130, 2),
+        ("mxfp8", 2, 17, 8, 4),
+        ("mxfp8", 2, 8, 17, 4),
+    ],
+)
+@pytest.mark.parametrize("layout", ["BSND", "BNSD"])
+def test_omni_calls_real_dense_public_api(monkeypatch, precision, layout, batch, q_len, kv_len, q_heads):
+    assert callable(getattr(mindiesd, "quant_attention", None)), "MindIE-SD must export quant_attention"
+    monkeypatch.setattr(flash_attn, "current_omni_platform", SimpleNamespace(is_npu=lambda: True, device_name="npu"))
+    monkeypatch.setattr(flash_attn, "get_current_diffusion_config_or_none", lambda: None)
+    quantized_inputs = []
+
+    def dynamic_quant(tensor, **kwargs):
+        quantized_inputs.append((tensor.clone(), kwargs))
+        shape = list(tensor.shape)
+        axis = kwargs.get("axis", -1)
+        shape[axis] = max(1, shape[axis] // 32)
+        return tensor, torch.ones(shape, dtype=torch.uint8)
+
+    monkeypatch.setattr(torch_npu, "npu_dynamic_block_quant", dynamic_quant, raising=False)
+    monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", dynamic_quant, raising=False)
+    execute = Mock(side_effect=lambda query, *args, **kwargs: (query, None))
+    metadata = Mock(return_value=torch.zeros(1, dtype=torch.int32))
+    if precision == "fp8":
+        monkeypatch.setattr(torch.ops.mindiesd, "fused_infer_attention_score_v2", execute, raising=False)
+    elif precision == "mxfp8":
+        monkeypatch.setattr(torch_npu, "npu_fused_infer_attention_score_v2", execute, raising=False)
+    else:
+        monkeypatch.setattr(torch.ops.mindiesd, "quant_flash_attn", execute, raising=False)
+        monkeypatch.setattr(torch.ops.mindiesd, "quant_flash_attn_metadata", metadata, raising=False)
+
+    query = torch.randn(batch, q_len, q_heads, 64, dtype=torch.bfloat16)
+    key = torch.randn(batch, kv_len, 2, 64, dtype=torch.bfloat16)
+    value = torch.randn_like(key)
+    if layout == "BNSD":
+        query, key, value = (tensor.transpose(1, 2) for tensor in (query, key, value))
+    impl = flash_attn.FlashAttentionImpl(
+        num_heads=q_heads,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout=layout,
+    )
+    output = impl.forward_npu(query, key, value, AttentionMetadata(extra={"kv_cache_dtype": precision}))
+    rotation = get_quant_attention_rotation(query.device, query.dtype, 64)
+    expected = query @ rotation
+    torch.testing.assert_close(output, expected)
+    assert output.shape == query.shape and len(quantized_inputs) == 3
+    execute.assert_called_once()
+    kwargs = execute.call_args.kwargs
+    assert kwargs["softmax_scale"] == 0.37
+    if precision == "fp8":
+        assert kwargs["input_layout"] == "BNSD"
+        assert quantized_inputs[0][0].shape == (2, 130, 64)
+        assert quantized_inputs[0][1]["row_block_size"] == 128
+    elif precision == "mxfp8":
+        args = execute.call_args.args
+        assert args[0].shape == (batch * q_len, q_heads, 64)
+        assert args[1].shape == args[2].shape == (batch * kv_len, 2, 64)
+        assert kwargs["input_layout"] == "TND"
+        assert kwargs["num_query_heads"] == q_heads and kwargs["num_key_value_heads"] == 2
+        assert kwargs["actual_seq_qlen"] == [q_len * (i + 1) for i in range(batch)]
+        assert kwargs["actual_seq_kvlen"] == [kv_len * (i + 1) for i in range(batch)]
+        assert [entry[1]["axis"] for entry in quantized_inputs] == [-1, -1, 0]
+    else:
+        seq_axis = 1 if layout == "BSND" else 2
+        assert quantized_inputs[0][0].shape[seq_axis] == 512
+        assert kwargs["seqused_q"].tolist() == kwargs["seqused_kv"].tolist() == [512]
+        assert kwargs["layout_q"] == kwargs["layout_out"] == layout
+        assert kwargs["q_dtype"] == torch_npu.float4_e2m1fn_x2
+        assert kwargs["q_descale_dtype"] == torch_npu.float8_e8m0fnu
+        assert execute.call_args.args[3].shape[:2] == (1, 2)
+        assert [entry[1]["axis"] for entry in quantized_inputs] == [-1, -1, seq_axis]
+        assert kwargs["metadata"] is metadata.return_value
+        metadata.assert_called_once()
+
+
+@pytest.mark.parametrize("precision", ["fp8", "mxfp4"])
+@pytest.mark.parametrize("length", [65, 600])
+def test_omni_calls_real_sparse_public_api(monkeypatch, precision, length):
+    rf = importlib.import_module("mindiesd.layers.flash_attn.sparse_flash_attn_rf_v3")
+    monkeypatch.setattr(flash_attn, "current_omni_platform", SimpleNamespace(is_npu=lambda: True, device_name="npu"))
+    for module in (flash_attn, rainfusion_attn):
+        monkeypatch.setattr(module, "get_current_diffusion_config_or_none", lambda: None)
+    rainfusion_attn._mindiesd_supports_precision.cache_clear()
+    monkeypatch.setattr(rf, "do_tensor_rearrange_only", lambda q, k, v, *a, **kw: (q, k, v))
+    monkeypatch.setattr(rf, "avgpool", lambda *a, **kw: torch.ones(1))
+    monkeypatch.setattr(rf, "_generate_mask_direct", lambda *a, **kw: torch.ones(1, 2, 8, 8))
+    monkeypatch.setattr(rf, "_bsa_inv_rearrange", lambda out, *a: out)
+
+    def quantize(tensor, **kwargs):
+        shape = list(tensor.shape)
+        axis = kwargs.get("axis", -1)
+        shape[axis] = max(1, shape[axis] // 32)
+        return tensor, torch.ones(shape, dtype=torch.uint8)
+
+    monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", quantize, raising=False)
+    monkeypatch.setattr(torch_npu, "npu_dynamic_block_quant", quantize, raising=False)
+    execute = Mock(side_effect=lambda **kw: (torch.ones_like(kw["query"]), None))
+    monkeypatch.setattr(torch.ops.mindiesd, "block_sparse_attention", execute, raising=False)
+    query = torch.randn(1, length, 2, 64, dtype=torch.bfloat16)
+    impl = rainfusion_attn.RainFusionAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout="BSND",
+        backend_kwargs={"sparsity": 0.8},
+    )
+    plan = rainfusion_attn.RainFusionPlan(used_len=length, prefix_len=0, latent_shape=(1, 1, length))
+    try:
+        output = impl._forward_sparse_npu(
+            query,
+            query,
+            query,
+            plan,
+            AttentionMetadata(extra={"kv_cache_dtype": precision}),
+        )
+    finally:
+        rainfusion_attn._mindiesd_supports_precision.cache_clear()
+    torch.testing.assert_close(output, torch.ones_like(query))
+    execute.assert_called_once()
+    kwargs = execute.call_args.kwargs
+    assert kwargs["actual_seq_lengths"] == kwargs["actual_seq_lengths_kv"] == [length]
+    assert kwargs["q_input_layout"] == "BNSD"
+    assert kwargs["block_sparse_mask"] is not None
+    if precision == "mxfp4":
+        assert kwargs["query"].shape[2] == (length + 63) // 64 * 64
+        assert kwargs["quant_mode"] == 2
+        assert kwargs["q_dtype"] == torch_npu.float4_e2m1fn_x2
+        assert kwargs["v_scale_dtype"] == torch_npu.float8_e8m0fnu

tests/diffusion/attention/test_mindie_runtime.py

@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU contract tests for MindIE adapters; no NPU kernels are mocked in production."""
+
+import sys
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention import layer as layer_mod
+from vllm_omni.diffusion.attention.backends import flash_attn, rainfusion_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+_DEFAULT_SKIP_STEPS = frozenset({0, 1, 38, 39})
+
+
+@pytest.fixture
+def runtime(monkeypatch):
+    mod = ModuleType("mindiesd")
+
+    def execute(q, k, v, *, precision, layout, scale, q_rot=None, k_rot=None):
+        return q
+
+    setattr(mod, "quant_attention", Mock(side_effect=execute))
+    monkeypatch.setitem(sys.modules, mod.__name__, mod)
+    for module in (flash_attn, layer_mod):
+        monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(is_npu=lambda: True, device_name="npu"))
+    monkeypatch.setattr(flash_attn, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(rainfusion_attn, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(rainfusion_attn, "is_forward_context_available", lambda: False)
+    supports_precision = rainfusion_attn._mindiesd_supports_precision
+    supports_precision.cache_clear()
+    yield mod
+    supports_precision.cache_clear()
+
+
+def flash(layout="BSND"):
+    return flash_attn.FlashAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout=layout,
+    )
+
+
+def sparse(**kwargs):
+    return rainfusion_attn.RainFusionAttentionImpl(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.37,
+        qkv_layout="BSND",
+        backend_kwargs={"sparsity": 0.8, **kwargs},
+    )
+
+
+def video_metadata(**extra):
+    return AttentionMetadata(
+        video_layout=VideoTokenLayout(prefix_len=0, latent_grid=(4, 32, 32), used_len=4096),
+        extra={"max_seqlen_q": 4096, **extra},
+    )
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp8", "mxfp4"])
+@pytest.mark.parametrize("layout", [None, "BSND", "BNSD"])
+def test_exact_runtime_layout_scale_and_rotation(runtime, method, layout):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    if layout == "BNSD":
+        q = q.transpose(1, 2)
+    impl = flash(layout=layout)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": method})
+    assert impl.forward_npu(q, q, q, metadata) is q
+    args, kwargs = runtime.quant_attention.call_args
+    assert args[0] is q  # no unconditional transpose or extra quantization
+    assert kwargs["layout"] == (layout or "BSND") and kwargs["scale"] == 0.37 and kwargs["precision"] == method
+    assert "attn_mask" not in kwargs
+    from vllm_omni.platforms.npu.quant.kv_quant_npu import get_quant_attention_rotation
+
+    assert kwargs["q_rot"] is kwargs["k_rot"]
+    torch.testing.assert_close(kwargs["q_rot"], get_quant_attention_rotation(q.device, q.dtype, 64))
+
+
+def test_missing_runtime_requires_config_change(runtime):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    del runtime.quant_attention
+    with pytest.raises(ImportError, match="diffusion_kv_cache_dtype"):
+        flash().forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": "mxfp4"}))
+
+
+def test_operator_failure_never_falls_back(runtime):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    runtime.quant_attention.side_effect = RuntimeError("operator failure")
+    with pytest.raises(RuntimeError, match="operator failure"):
+        flash().forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": "mxfp4"}))
+    runtime.quant_attention.assert_called_once()
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp8", "mxfp4"])
+def test_unsupported_shape_requires_config_change(runtime, method):
+    q = torch.randn(1, 128, 2, 96, dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match="diffusion_kv_cache_dtype"):
+        flash().forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": method}))
+    runtime.quant_attention.assert_not_called()
+
+
+def test_packed_metadata_never_reaches_quant_runtime(runtime):
+    q = torch.randn(1, 128, 2, 64, dtype=torch.bfloat16)
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": "mxfp8", "cu_seqlens_q": torch.tensor([0, 64, 128])})
+    with pytest.raises(ValueError, match="packed/varlen"):
+        flash().forward_npu(q, q, q, metadata)
+    runtime.quant_attention.assert_not_called()
+
+
+def make_layer(runtime, monkeypatch, method="mxfp8", skip_steps=_DEFAULT_SKIP_STEPS):
+    layer = object.__new__(layer_mod.Attention)
+    torch.nn.Module.__init__(layer)
+    layer.attention = flash()
+    layer.attn_backend = flash_attn.FlashAttentionBackend
+    layer._disable_kv_quant = False
+    layer.layer_idx = 3
+    cfg = SimpleNamespace(
+        diffusion_kv_cache_dtype=method,
+        diffusion_kv_cache_skip_step_indices=skip_steps,
+        diffusion_kv_cache_skip_layer_indices=None,
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    layer._init_kv_cache_quantization(cfg)
+    return layer
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp8", "mxfp4"])
+def test_global_quant_method_is_published_per_forward(runtime, monkeypatch, method):
+    layer = make_layer(runtime, monkeypatch, method=method, skip_steps=None)
+    assert layer._with_kv_cache_dtype(None).extra["kv_cache_dtype"] == method
+
+
+@pytest.mark.parametrize("method", [None, "auto", "float"])
+def test_optout_without_quant_policy_keeps_metadata_unchanged(runtime, monkeypatch, method):
+    layer = make_layer(runtime, monkeypatch)
+    layer._disable_kv_quant = True
+    layer._init_kv_cache_quantization(
+        SimpleNamespace(
+            diffusion_kv_cache_dtype=method,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+            parallel_config=SimpleNamespace(ring_degree=1),
+        )
+    )
+    source = video_metadata()
+    assert layer._with_kv_cache_dtype(source) is source
+
+
+def test_step_skip_is_request_local_and_cross_optout(runtime, monkeypatch):
+    layer = make_layer(runtime, monkeypatch)
+    ctx = SimpleNamespace(denoise_step_idx=0)
+    monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(layer_mod, "get_forward_context", lambda: ctx)
+    source = video_metadata()
+    for step in list(range(40)) * 2:
+        ctx.denoise_step_idx = step
+        resolved = layer._with_kv_cache_dtype(source)
+        assert resolved.extra["kv_cache_dtype"] == ("float" if step in (0, 1, 38, 39) else "mxfp8")
+        assert resolved.video_layout is source.video_layout
+    assert source.extra == {"max_seqlen_q": 4096}
+    layer._disable_kv_quant = True
+    assert layer._with_kv_cache_dtype(source).extra["kv_cache_dtype"] == "float"
+
+
+def test_ring_attention_is_rejected(runtime, monkeypatch):
+    layer = make_layer(runtime, monkeypatch)
+    cfg = SimpleNamespace(diffusion_kv_cache_dtype="mxfp8", parallel_config=SimpleNamespace(ring_degree=2))
+    with pytest.raises(ValueError, match="ring"):
+        layer._init_kv_cache_quantization(cfg)
+
+
+def test_sparse_tail_steps_and_dense_quant_dispatch(runtime, monkeypatch):
+    ctx = SimpleNamespace(denoise_step_idx=49, total_denoise_steps=50)
+    monkeypatch.setattr(rainfusion_attn, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(rainfusion_attn, "get_forward_context", lambda: ctx)
+    impl = sparse(start_step=2, end_step=2)
+    metadata = video_metadata(kv_cache_dtype="fp8")
+    assert impl._resolve_plan(metadata) is None
+    q = torch.randn(1, 4096, 2, 64, dtype=torch.bfloat16)
+    impl.forward_npu(q, q, q, metadata)
+    assert runtime.quant_attention.call_args.kwargs["precision"] == "fp8"
+
+
+@pytest.mark.parametrize("method", ["fp8", "mxfp4"])
+def test_sparse_quant_uses_public_runtime(runtime, monkeypatch, method):
+    runtime.sparse_attention = Mock(side_effect=lambda q, k, v, **kwargs: q)
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
+    q = torch.randn(1, 4224, 2, 64, dtype=torch.bfloat16)
+    out = sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype=method))
+    assert out.shape == q.shape and not out[:, 4096:].any()
+    kwargs = runtime.sparse_attention.call_args.kwargs
+    assert (kwargs["precision"], kwargs["sparse_type"], kwargs["inner_precise"]) == (method, "rf_v3", 4)
+    runtime.quant_attention.assert_not_called()
+
+
+def test_sparse_custom_mask_stays_dense(runtime):
+    metadata = AttentionMetadata(
+        attn_mask=torch.ones(1, 4096, dtype=torch.bool),
+        video_layout=VideoTokenLayout(prefix_len=0, latent_grid=(4, 32, 32)),
+        extra={"max_seqlen_q": 4096},
+    )
+    assert sparse()._resolve_plan(metadata) is None
+
+
+def test_sparse_model_padding_mask_is_equivalent_to_trimming(runtime):
+    metadata = video_metadata()
+    metadata.attn_mask = torch.arange(4224)[None] < 4096
+    assert sparse()._resolve_plan(metadata).used_len == 4096
+
+
+def test_bsa_operator_failure_never_retries(runtime, monkeypatch):
+    runtime.sparse_attention = Mock(side_effect=RuntimeError("BSA operator failure"))
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
+    q = torch.randn(1, 4096, 2, 64, dtype=torch.bfloat16)
+    with pytest.raises(RuntimeError, match="BSA operator failure"):
+        sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype="mxfp4"))
+    runtime.sparse_attention.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "method,shape,supports_precision,error",
+    [
+        ("mxfp8", (1, 4096, 2, 64), True, "sparse mxfp8 is not supported"),
+        ("mxfp4", (1, 4096, 2, 64), False, "explicitly support precision"),
+        ("fp8", (2, 4096, 2, 64), True, "batch size 1"),
+        ("fp8", (1, 4096, 2, 32), True, "power-of-two head dimension"),
+        ("mxfp4", (1, 4096, 2, 96), True, "power-of-two head dimension"),
+    ],
+)
+def test_invalid_sparse_quant_requires_config_change(runtime, monkeypatch, method, shape, supports_precision, error):
+    runtime.sparse_attention = Mock(side_effect=lambda q, k, v, **kwargs: q)
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: supports_precision)
+    q = torch.randn(shape, dtype=torch.bfloat16)
+    with pytest.raises(ValueError, match=error):
+        sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype=method))
+    runtime.sparse_attention.assert_not_called()
+
+
+def test_layer_selector_requires_index_except_cross_optout(runtime, monkeypatch):
+    layer = make_layer(runtime, monkeypatch)
+    layer.layer_idx = None
+    config = SimpleNamespace(
+        diffusion_kv_cache_dtype="mxfp8",
+        diffusion_kv_cache_skip_step_indices=None,
+        diffusion_kv_cache_skip_layer_indices={0, 3},
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    with pytest.raises(ValueError, match="parseable transformer block index"):
+        layer._init_kv_cache_quantization(config)
+    layer._disable_kv_quant = True
+    layer._init_kv_cache_quantization(config)
+    assert layer._with_kv_cache_dtype(None).extra["kv_cache_dtype"] == "float"
+
+
+@pytest.mark.parametrize("prefix", ["transformer.blocks.3.attn1", "transformer_2.blocks.3.attn1"])
+@pytest.mark.parametrize("method", ["fp8", "mxfp4"])
+def test_expert_layer_and_step_skip_remain_sparse(runtime, monkeypatch, prefix, method):
+    def execute(q, k, v, *, precision="bf16", **kwargs):
+        return q
+
+    runtime.sparse_attention = Mock(wraps=execute)
+    monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
+    layer = make_layer(runtime, monkeypatch)
+    layer.layer_idx = rainfusion_attn._try_extract_layer_index(prefix)
+    assert layer.layer_idx == 3
+    config = SimpleNamespace(
+        diffusion_kv_cache_dtype=method,
+        diffusion_kv_cache_skip_step_indices={0, 1, 38, 39},
+        diffusion_kv_cache_skip_layer_indices={3},
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    impl = sparse()
+    layer.attention = impl
+    layer.attn_backend = rainfusion_attn.RainFusionAttentionBackend
+    layer._init_kv_cache_quantization(config)
+    ctx = SimpleNamespace(denoise_step_idx=0)
+    monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(layer_mod, "get_forward_context", lambda: ctx)
+    source = video_metadata(kv_cache_dtype="mxfp8")
+    q = torch.randn(1, 4096, 2, 64, dtype=torch.bfloat16)
+    for layer_idx, step in ((3, 0), (3, 2), (4, 0), (4, 1), (4, 2), (4, 37), (4, 38), (4, 39), (4, 0)):
+        expected = "bf16" if layer_idx == 3 or step in (0, 1, 38, 39) else method
+        layer.layer_idx = layer_idx
+        ctx.denoise_step_idx = step
+        metadata = layer._with_kv_cache_dtype(source)
+        impl.forward_npu(q, q, q, metadata)
+        assert runtime.sparse_attention.call_args.kwargs["precision"] == expected
+    runtime.quant_attention.assert_not_called()
+    assert source.extra["kv_cache_dtype"] == "mxfp8"
+
+
+def test_custom_attention_initialization_keeps_upstream_optout(monkeypatch):
+    monkeypatch.setattr(layer_mod, "get_current_diffusion_config_or_none", lambda: None)
+    monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kw: object())
+    monkeypatch.setattr(layer_mod, "NoParallelAttention", lambda: object())
+    custom = torch.nn.Identity()
+    layer = layer_mod.Attention(
+        num_heads=2,
+        head_size=64,
+        softmax_scale=0.125,
+        causal=False,
+        custom_attention=custom,
+        skip_sequence_parallel=True,
+    )
+    assert layer.attention is custom and layer._kv_cache_dtype is None
+
+
+def test_legacy_global_fp8_uses_implicit_bsnd(runtime):
+    q = torch.randn(1, 17, 2, 64, dtype=torch.bfloat16)
+    impl = flash(layout=None)
+    assert impl.forward_npu(q, q, q, AttentionMetadata(extra={"kv_cache_dtype": "fp8"})) is q
+    assert runtime.quant_attention.call_args.kwargs["layout"] == "BSND"
+    assert runtime.quant_attention.call_args.kwargs["precision"] == "fp8"

tests/diffusion/attention/test_rainfusion_plan.py

@@ -133,18 +133,17 @@ def test_invalid_ref2va_video_spans_fall_back_to_dense():
     assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
 
 
-def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
+def test_validate_available_accepts_single_video_mindiesd(monkeypatch):
     mindiesd = types.ModuleType("mindiesd")
 
     def sparse_attention(query, key, value, **kwargs):
         return query
 
-    mindiesd.sparse_attention = sparse_attention
+    monkeypatch.setattr(mindiesd, "sparse_attention", sparse_attention, raising=False)
     monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
     monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
 
-    with pytest.raises(ValueError, match="video_spans"):
-        RainFusionAttentionBackend.validate_available()
+    RainFusionAttentionBackend.validate_available()
 
 
 def test_validate_available_accepts_new_mindiesd(monkeypatch):
@@ -153,7 +152,7 @@ def test_validate_available_accepts_new_mindiesd(monkeypatch):
     def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
         return query
 
-    mindiesd.sparse_attention = sparse_attention
+    monkeypatch.setattr(mindiesd, "sparse_attention", sparse_attention, raising=False)
     monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
     monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
 
@@ -300,22 +299,21 @@ def _fake_mindiesd_module():
     import types
 
     fake = types.ModuleType("mindiesd")
-    fake.sparse_attention = lambda *args, **kwargs: None
+    setattr(fake, "sparse_attention", lambda *args, **kwargs: None)
     return fake
 
 
 def test_precision_non_bf16_requires_mindiesd_support():
-    """precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError."""
+    """Unsupported sparse precision must fail before any operator executes."""
     import sys
 
     impl = make_impl(precision="mix")
-    sys.modules["mindiesd"] = _fake_mindiesd_module()
-    try:
+    with mock.patch.dict(sys.modules, {"mindiesd": _fake_mindiesd_module()}):
         with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False):
-            with pytest.raises(RuntimeError, match="requires MindIE-SD"):
-                impl._forward_sparse_npu(None, None, None, None)
-    finally:
-        sys.modules.pop("mindiesd", None)
+            with pytest.raises(ValueError, match="explicitly support precision"):
+                impl._forward_sparse_npu(
+                    None, None, None, RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2])
+                )
 
 
 def test_precision_non_bf16_passes_gate_when_supported():
@@ -323,8 +321,7 @@ def test_precision_non_bf16_passes_gate_when_supported():
     import sys
 
     impl = make_impl(precision="mix")
-    sys.modules["mindiesd"] = _fake_mindiesd_module()
-    try:
+    with mock.patch.dict(sys.modules, {"mindiesd": _fake_mindiesd_module()}):
         with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True):
             # q/k/v shapes: [B, S, N, D]; plan geometry must match S.
             q = torch.randn(1, 8, 4, 128)
@@ -332,5 +329,3 @@ def test_precision_non_bf16_passes_gate_when_supported():
             # The gate must pass; the fake mindiesd returns None so no crash.
             out = impl._forward_sparse_npu(q, q, q, plan)
             assert out is None
-    finally:
-        sys.modules.pop("mindiesd", None)

tests/diffusion/models/wan2_2/test_wan22_mindie_metadata.py

@@ -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