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

docs/user_guide/diffusion/attention_backends/rainfusion.md

@@ -78,3 +78,35 @@ Protected video tails require a compatible MindIE-SD release.
 For common configuration and selector behavior, see the
 [attention backend overview](../attention_backends.md) and the
 [backend selection design](../../../design/feature/attention_backend_selection.md).
+
+## Wan2.2 T2V quantization on Ascend
+
+`RAINFUSION_ATTN` accepts `diffusion_kv_cache_dtype: fp8` or `mxfp4` for
+Wan2.2 T2V A14B.
+MindIE-SD must expose `sparse_attention` with `rf_v3` and an explicit `precision`
+argument. Omni passes the selected precision directly to this API; the installed
+native operators must support that precision. Unsupported inputs or an
+incompatible Python API raise an actionable error; Omni does not change BSA
+precision automatically. Update `diffusion_kv_cache_dtype` to a precision
+supported by the installed MindIE-SD, or unset it. Native execution errors also
+propagate without retry.
+BSA FP8 and MXFP4 retain MindIE-SD's default rotation policy.
+
+`diffusion_kv_cache_skip_layers`/`diffusion_kv_cache_skip_steps` explicitly
+select floating-point BSA and keep attention sparse. The existing `block_sparse`
+warmup/layer exclusions and short-sequence threshold instead select Dense
+attention, using the configured Dense precision. Leave `block_sparse.precision`
+at its default when using `diffusion_kv_cache_dtype` as the shared Dense/BSA
+precision setting.
+
+Configure BSA quantization in the model's existing deployment YAML with
+`backend: RAINFUSION_ATTN` and stage-level `diffusion_kv_cache_dtype: fp8` or
+`mxfp4`. The precision and skip selectors are siblings of
+`diffusion_attention_config`; they are not per-role `quant` fields.
+For 40 denoising steps, `diffusion_kv_cache_skip_steps: "0,1,38,39"` and
+`diffusion_kv_cache_skip_layers: "0,39"` are recommended starting points. These
+fixed step indices do not automatically follow a changed step count. Selected
+forwards remain floating-point BSA; unsupported BSA precision is an error rather
+than an automatic fallback. See
+[quantized attention](../../quantization/quantized_kvcache.md#wan22-t2v-quantized-attention-on-ascend)
+for the shared configuration structure and index semantics.

docs/user_guide/quantization/overview.md

@@ -13,7 +13,7 @@ For the internal architecture and backend extension points, see the
 | Mode | Guide | Description | Methods |
 | ------ | ------- | ------------- | --------- |
 | Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, BitsAndBytes W4, MXFP8 W8A8, MXFP4 W4A4 |
-| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
+| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion attention tensors during inference. | FP8, MXFP8, MXFP4 |
 | Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, TorchAO, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
 
 ## Hardware Support

docs/user_guide/quantization/quantized_kvcache.md

@@ -2,41 +2,36 @@
 
 ## Overview
 
-In DiT-based image and video generation, Flash Attention can take a large share
-of denoising time, especially for high-resolution or long-frame workloads.
-vLLM-Omni supports online FP8 quantization for eligible diffusion Flash
-Attention (FA) to reduce FA latency while keeping model weights in their
-original dtype.
+In DiT-based image and video generation, attention can take a large share of
+denoising time, especially for high-resolution or long-frame workloads.
+vLLM-Omni supports online Q/K/V quantization for eligible diffusion attention
+paths while keeping model weights in their original dtype.
 
-This feature is configured through `diffusion_kv_cache_dtype` on
-`OmniDiffusionConfig` (CLI: `--diffusion-kv-cache-dtype`). It is intentionally
-**not** the same as vLLM's `--kv-cache-dtype`, which controls autoregressive
-language-model KV cache storage and defaults to `"auto"`. Diffusion FA
-quantization uses the dedicated diffusion flags so omni serve does not inherit
-that default.
+Runtime attention precision is configured through `diffusion_kv_cache_dtype`
+on `OmniDiffusionConfig` (CLI: `--diffusion-kv-cache-dtype`). This setting is
+separate from vLLM's `--kv-cache-dtype`, which controls autoregressive
+language-model KV cache storage.
 
 In vLLM-Omni diffusion pipelines, this is a runtime FA path: Q/K/V tensors are
 dynamically quantized before the attention operator. It does not quantize model
 weights and is separate from [FP8 W8A8](fp8.md), [Int8 W8A8](int8.md), or
 pre-quantized checkpoint formats.
 
-If `diffusion_kv_cache_dtype` is not set, behavior is unchanged and attention
-runs in the native dtype.
+If `diffusion_kv_cache_dtype` is not set, attention runs in the native dtype.
 
 ## Hardware Support
 
-| Device | FP8 FA |
-|--------|--------|
-| Ascend NPU | ✅ |
-| NVIDIA GPU | ❌ |
-| AMD ROCm | ❌ |
-| Intel XPU | ❌ |
+| Device | Dense FP8 | Dense MXFP8/MXFP4 | BSA FP8/MXFP4 |
+|--------|------------|-------------------|----------------|
+| Ascend NPU | ✅ | ✅ | ✅ |
+| NVIDIA GPU | ❌ | ❌ | ❌ |
+| AMD ROCm | ❌ | ❌ | ❌ |
+| Intel XPU | ❌ | ❌ | ❌ |
 
 Legend: `✅` supported, `❌` unsupported.
 
-FP8 FA is currently implemented only for the NPU Flash Attention backend. Other
-backends do not support `diffusion_kv_cache_dtype="fp8"` for diffusion attention
-and fall back to native dtype execution.
+These runtime methods are implemented only by the NPU Flash Attention and BSA
+backends. Other backends reject an incompatible explicit configuration.
 
 ## Model Type Support
 
@@ -44,7 +39,7 @@ and fall back to native dtype execution.
 
 | Model | Scope | Status | Notes |
 |-------|-------|--------|-------|
-| Wan2.2 | Eligible DiT full-attention FA on Ascend NPU | Tested | Compare quality and latency against a BF16 baseline before production use |
+| Wan2.2 | Self-attention on Ascend NPU | NPU interface and T2V smoke tested | Compare quality and latency against a BF16 baseline before production use |
 | Other diffusion models | Eligible DiT full-attention FA on Ascend NPU | Not tested | You can try `diffusion_kv_cache_dtype="fp8"`; tune `diffusion_kv_cache_skip_steps` and `diffusion_kv_cache_skip_layers` when higher precision is needed |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -105,12 +100,13 @@ The legacy keyword aliases `kv_cache_dtype`, `kv_cache_skip_steps`, and
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `diffusion_kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends |
+| `diffusion_kv_cache_dtype` | str \| None | `None` | Runtime attention method. Dense FA supports `fp8`, `mxfp8`, and `mxfp4`; BSA supports `fp8` and `mxfp4` |
 | `diffusion_kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` |
 | `diffusion_kv_cache_skip_layers` | str \| None | `None` | Transformer layer selector to keep in native dtype, for example `"0-2,10"` |
 
 Selectors use comma-separated integers and inclusive ranges. Listed steps or
-layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path.
+layers run floating-point attention through the selected Dense or BSA backend;
+all other eligible forwards use `diffusion_kv_cache_dtype`.
 
 ## Validation and Notes
 
@@ -123,3 +119,60 @@ layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path.
 4. Report both latency and quality results when enabling this option for a new
    model. For image or video models, include visual comparison and quantitative
    metrics when available, such as PSNR or SSIM.
+
+## Wan2.2 T2V quantized attention on Ascend
+
+Wan2.2 T2V A14B supports quantized self-attention through MindIE-SD. Model
+weights are unchanged, and model-level opt-out keeps cross-attention in its
+native dtype.
+
+| Attention path | self-attention `backend` | Supported `diffusion_kv_cache_dtype` |
+| --- | --- | --- |
+| Dense FA | `FLASH_ATTN` | `fp8`, `mxfp8`, `mxfp4` |
+| BSA | `RAINFUSION_ATTN` | `fp8`, `mxfp4` |
+
+Configure the method and fallback selectors once at the diffusion stage. The
+per-role attention config only selects the self-attention backend:
+
+```yaml
+stages:
+  - stage_id: 0
+    diffusion_kv_cache_dtype: "mxfp8"
+    diffusion_kv_cache_skip_layers: "0,39"
+    diffusion_kv_cache_skip_steps: "0,1,38,39"
+    diffusion_attention_config:
+      per_role:
+        self:
+          backend: FLASH_ATTN  # Use RAINFUSION_ATTN with fp8 or mxfp4 for BSA.
+```
+
+For 40 denoising steps, `diffusion_kv_cache_skip_steps: "0,1,38,39"` and
+`diffusion_kv_cache_skip_layers: "0,39"` are recommended starting points. Steps
+are zero-based across the complete request and do not reset when Wan switches
+transformers; layer indices are zero-based and local to each transformer.
+Selectors accept comma-separated indices and inclusive ranges such as
+`"0,3-5"`. A selected forward uses floating-point attention while preserving
+the Dense or BSA path.
+
+Unsupported precision/input combinations raise an error and must be corrected
+in the configuration; operator errors are not retried with another precision.
+
+For BSA, change the backend to `RAINFUSION_ATTN` and set
+`diffusion_kv_cache_dtype` to `fp8` or `mxfp4`.
+See [RainFusion attention](../diffusion/attention_backends/rainfusion.md) for
+sparse-path configuration and behavior.
+
+### Requirement and interface test
+
+Until these APIs are available in a MindIE-SD release, use revision
+[`8637b5333b0225381b215390fd09a8732e671cc4`](https://gitcode.com/zqxu/MindIE-SD/commit/8637b5333b0225381b215390fd09a8732e671cc4).
+Build the Python package, PyTorch plugin, and custom operators from that same
+checkout. With CANN sourced, run the NPU interface tests from the Omni checkout:
+
+```bash
+python -m pytest tests/platforms/npu/quant/test_kv_quant_npu.py \
+    -k real_npu -vv -s -o addopts=''
+```
+
+A skipped NPU test is not a pass. Validate generated video quality against a
+floating-point run with the same prompt, seed, resolution, and step count.

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

tests/platforms/npu/quant/test_kv_quant_npu.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for NPU FP8 KV quantization helpers.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Tests for seeded rotations and the quantized NPU attention backend.
 
 These tests load ``kv_quant_npu`` from its source file via ``importlib`` so
 the test module itself does not ``import vllm_omni`` (which would pull
@@ -10,10 +10,8 @@ the test module itself does not ``import vllm_omni`` (which would pull
 from __future__ import annotations
 
 import importlib.util
-import math
 from pathlib import Path
-from types import ModuleType, SimpleNamespace
-from typing import Any
+from types import ModuleType
 
 import pytest
 import torch
@@ -61,147 +59,39 @@ def _npu_smoke_available() -> bool:
 npu_smoke = pytest.mark.skipif(not _npu_smoke_available(), reason="NPU device or torch_npu not available.")
 
 
-def test_is_quantized_kv_cache() -> None:
-    assert kv_quant_npu.is_quantized_kv_cache("fp8")
-    assert not kv_quant_npu.is_quantized_kv_cache(None)
-    assert not kv_quant_npu.is_quantized_kv_cache("int8")
-
-
-class TestKVQuantNPUUnit:
-    @pytest.fixture(autouse=True)
-    def clear_rot_cache(self):
-        kv_quant_npu._ROT_MATRIXS.clear()
-
-    def test_get_rot_matrix_caches_by_device_dtype_and_head_dim(self) -> None:
-        calls = {"count": 0}
-
-        class FakeQuaRotMode:
-            HADAMARD = "hadamard"
-
-        def fake_create_rot(mode, head_dim, seed):
-            calls["count"] += 1
-            assert mode == FakeQuaRotMode.HADAMARD
-            assert seed == 425500
-            return torch.eye(head_dim, dtype=torch.float32)
-
-        device = torch.device("cpu")
-        rot_1 = kv_quant_npu._get_rot_matrix(device, torch.float16, 8, FakeQuaRotMode, fake_create_rot)
-        rot_2 = kv_quant_npu._get_rot_matrix(device, torch.float16, 8, FakeQuaRotMode, fake_create_rot)
-        rot_3 = kv_quant_npu._get_rot_matrix(device, torch.bfloat16, 8, FakeQuaRotMode, fake_create_rot)
-        rot_4 = kv_quant_npu._get_rot_matrix(device, torch.float16, 16, FakeQuaRotMode, fake_create_rot)
-
-        assert calls["count"] == 3
-        assert rot_1 is rot_2
-        assert rot_3.dtype == torch.bfloat16
-        assert rot_4.shape == (16, 16)
-
-    @pytest.fixture
-    def fake_quant_ops(self, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]:
-        captured: dict[str, Any] = {
-            "fa_calls": [],
-            "npu_kwargs": None,
-            "out_shape": None,
-        }
-
-        class FakeTorchNPU:
-            float8_e4m3fn = "fp8_marker"
-
-            @staticmethod
-            def npu_fused_infer_attention_score_v2(q, k, v, **kwargs):
-                del q, k, v
-                captured["npu_kwargs"] = kwargs
-                out_shape = captured["out_shape"]
-                return (torch.ones(out_shape, dtype=torch.float32),)
-
-        def fake_fa_block_quant_preprocess(x, block_size, dst_type, layout):
-            captured["fa_calls"].append(
-                {
-                    "block_size": block_size,
-                    "layout": layout,
-                    "dst_type": dst_type,
-                    "shape": tuple(x.shape),
-                }
-            )
-            scale = torch.full((1,), float(block_size), dtype=torch.float32)
-            return x, scale
-
-        fake_qua_rot_mode = SimpleNamespace(HADAMARD="hadamard")
-
-        def fake_create_rot(mode, head_dim, seed):
-            assert mode == "hadamard"
-            assert seed == 425500
-            return torch.eye(head_dim, dtype=torch.float32)
-
-        monkeypatch.setattr(
-            kv_quant_npu,
-            "_load_quant_ops",
-            lambda: (FakeTorchNPU, fake_fa_block_quant_preprocess, fake_qua_rot_mode, fake_create_rot),
-        )
-
-        return captured
-
-    @staticmethod
-    def _make_qkv(shape: tuple[int, int, int, int]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        query = torch.randn(*shape, dtype=torch.float32)
-        key = torch.randn(*shape, dtype=torch.float32)
-        value = torch.randn(*shape, dtype=torch.float32)
-        return query, key, value
-
-    @pytest.mark.parametrize(
-        "layout,input_shape,out_shape,softmax_scale,expected_scale",
-        [
-            ("BNSD", (2, 3, 4, 8), (2, 3, 6, 8), None, 1.0 / math.sqrt(8)),
-            ("BSND", (2, 4, 3, 8), (2, 6, 3, 8), 0.125, 0.125),
-        ],
-    )
-    def test_fp8_rotate_quant_fa_layouts_scale_and_crop(
-        self,
-        fake_quant_ops: dict[str, Any],
-        layout: str,
-        input_shape: tuple[int, int, int, int],
-        out_shape: tuple[int, int, int, int],
-        softmax_scale: float | None,
-        expected_scale: float,
-    ) -> None:
-        query, key, value = self._make_qkv(input_shape)
-        fake_quant_ops["out_shape"] = out_shape
-
-        out = kv_quant_npu.fp8_rotate_quant_fa(query, key, value, layout=layout, softmax_scale=softmax_scale)
-
-        assert out.shape == query.shape
-        assert out.dtype == query.dtype
-        assert fake_quant_ops["npu_kwargs"]["input_layout"] == layout
-        # BNSD: shape[1]==heads, BSND: shape[2]==heads.
-        expected_heads = input_shape[1] if layout == "BNSD" else input_shape[2]
-        assert fake_quant_ops["npu_kwargs"]["num_query_heads"] == expected_heads
-        assert fake_quant_ops["npu_kwargs"]["softmax_scale"] == pytest.approx(expected_scale)
-        assert [call["block_size"] for call in fake_quant_ops["fa_calls"]] == [128, 256, 256]
-
-    def test_fp8_rotate_quant_fa_invalid_layout_raises(self, fake_quant_ops) -> None:
-        query = torch.randn(1, 2, 3, 4, dtype=torch.float32)
-        key = torch.randn(1, 2, 3, 4, dtype=torch.float32)
-        value = torch.randn(1, 2, 3, 4, dtype=torch.float32)
-        fake_quant_ops["out_shape"] = (1, 2, 3, 4)
-
-        with pytest.raises(ValueError, match="unsupported layout"):
-            kv_quant_npu.fp8_rotate_quant_fa(query, key, value, layout="INVALID")
+def test_rotation_is_fixed_orthogonal_cached_and_preserves_rng():
+    kv_quant_npu.get_quant_attention_rotation.cache_clear()
+    before = torch.random.get_rng_state()
+    rot = kv_quant_npu.get_quant_attention_rotation(torch.device("cpu"), torch.float32, 64)
+    assert rot is kv_quant_npu.get_quant_attention_rotation(torch.device("cpu"), torch.float32, 64)
+    torch.testing.assert_close(rot @ rot.T, torch.eye(64))
+    assert torch.equal(before, torch.random.get_rng_state())
 
 
 @npu_smoke
-class TestKVQuantNPUSmoke:
-    """Smoke tests using real torch_npu/mindiesd stack, only on NPU."""
+@pytest.mark.npu
+@pytest.mark.parametrize("method", ["fp8", "mxfp8", "mxfp4"])
+@pytest.mark.parametrize("layout", ["BSND", "BNSD"])
+@pytest.mark.parametrize("seq_len", [256, 75600])
+def test_dense_runtime_backend_real_npu(method, layout, seq_len):
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+    from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 
-    def test_fp8_rotate_quant_fa_real_npu_shape_contract(self):
-        try:
-            kv_quant_npu._load_quant_ops.cache_clear()
-            kv_quant_npu._load_quant_ops()
-        except ImportError:
-            pytest.skip("NPU quant dependencies are not fully installed.")
-
-        query = torch.randn(1, 2, 4, 64, dtype=torch.float16, device="npu")
-        key = torch.randn(1, 2, 4, 64, dtype=torch.float16, device="npu")
-        value = torch.randn(1, 2, 4, 64, dtype=torch.float16, device="npu")
-
-        out = kv_quant_npu.fp8_rotate_quant_fa(query, key, value, layout="BNSD")
-        assert out.shape == query.shape
-        assert out.dtype == query.dtype
+    # On NPU, missing Runtime symbols must fail this qualification test.
+    query = torch.randn(1, seq_len, 2, 128, dtype=torch.bfloat16, device="npu") * 0.1
+    # Distinct constant values per head detect head/layout mixing and tail writes.
+    value = torch.tensor([0.5, -0.5], dtype=query.dtype, device=query.device)
+    value = value.view(1, 1, 2, 1).expand_as(query).contiguous()
+    if layout == "BNSD":
+        query, value = query.transpose(1, 2), value.transpose(1, 2)
+    impl = FlashAttentionImpl(
+        num_heads=2,
+        head_size=128,
+        softmax_scale=128**-0.5,
+        qkv_layout=layout,
+    )
+    metadata = AttentionMetadata(extra={"kv_cache_dtype": method})
+    out = impl.forward_npu(query, query, value, metadata)
+    assert out.shape == query.shape and out.dtype == query.dtype
+    assert torch.isfinite(out).all()
+    torch.testing.assert_close(out.float().cpu(), value.float().cpu(), rtol=0.02, atol=0.02)

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import os
+from collections.abc import Callable
 from functools import cache, partial
 
 import torch
@@ -75,18 +76,18 @@ class FlashAttentionBackend(AttentionBackend):
 
 
 class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
-    # Per-platform FP8 KV quantization support.
-    # To enable FP8 on a new platform, add its OmniPlatformEnum value here
+    # Per-platform attention quantization support.
+    # To enable a method on a new platform, add its OmniPlatformEnum value here
     # and handle kv_cache_dtype in the corresponding forward_{platform}().
     #
-    # TODO(quant-backend): The FP8 quant path currently lives inside
+    # TODO(quant-backend): The quantized path currently lives inside
     # FlashAttentionImpl gated by ``attn_metadata.extra["kv_cache_dtype"]``.
     # Eventually extract it into a dedicated FlashAttentionQuantBackend so
-    # backend selection (not metadata) decides quant. Until then, model
-    # authors can opt a specific Attention layer out via
+    # backend selection decides quant.
+    # Until then, model authors can opt a specific Attention layer out via
     # ``Attention(disable_kv_quant=True)``.
     _supported_kv_cache_dtypes = {
-        "npu": {"fp8"},
+        "npu": {"fp8", "mxfp8", "mxfp4"},
     }
 
     def __init__(
@@ -526,11 +527,47 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
     ) -> torch.Tensor:
         """NPU attention implementation using mindiesd."""
 
-        kv_cache_dtype = attn_metadata.extra.get("kv_cache_dtype") if attn_metadata else None
-        if kv_cache_dtype is not None:
+        extra = attn_metadata.extra if attn_metadata else {}
+        method = extra.get("kv_cache_dtype")
+        if method not in (None, "float", "auto"):
             return self.forward_fa_quant_npu(query, key, value, attn_metadata)
         return self.forward_fa_npu(query, key, value, attn_metadata)
 
+    @staticmethod
+    def _load_quant_runtime(method: str) -> Callable[..., torch.Tensor]:
+        if method not in ("fp8", "mxfp8", "mxfp4"):
+            raise ValueError(f"Unsupported NPU attention quantization method {method!r}.")
+        from mindiesd import quant_attention
+
+        return quant_attention
+
+    def _validate_quant_request(
+        self,
+        method: str,
+        query: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> None:
+        extra = attn_metadata.extra if attn_metadata else {}
+        reason = None
+        if self.causal:
+            reason = "causal quantized FA is not supported"
+        elif any(name in extra for name in ("cu_seqlens_q", "cu_seqlens_k")) or extra.get("npu_attn_varlen"):
+            reason = "packed/varlen metadata requires the float attention path"
+        elif attn_metadata is not None and attn_metadata.full_attn_spans is not None:
+            reason = "piecewise attention requires the float attention path"
+        elif attn_metadata is not None and attn_metadata.attn_mask is not None:
+            reason = "caller masks require the float attention path"
+        elif method in ("fp8", "mxfp8", "mxfp4") and (
+            query.ndim != 4 or query.shape[-1] <= 0 or query.shape[-1] & (query.shape[-1] - 1)
+        ):
+            reason = "generated Hadamard rotations require a four-dimensional input and power-of-two head size"
+        if reason is not None:
+            raise ValueError(
+                f"NPU attention precision {method!r} is unavailable: {reason}. "
+                "Set diffusion_kv_cache_dtype to a supported method, or set it to "
+                "'auto'. Automatic precision fallback is not performed."
+            )
+
     def forward_fa_quant_npu(
         self,
         query: torch.Tensor,
@@ -538,18 +575,29 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
-        from vllm_omni.platforms.npu.quant.kv_quant_npu import fp8_rotate_quant_fa
+        extra = attn_metadata.extra if attn_metadata else {}
+        method = extra.get("kv_cache_dtype", "fp8")
+        layout = self.qkv_layout or "BSND"
+        self._validate_quant_request(method, query, attn_metadata)
+        try:
+            runtime = self._load_quant_runtime(method)
+        except ImportError as exc:
+            raise ImportError(
+                f"NPU attention precision {method!r} requires a compatible MindIE-SD "
+                "quant Runtime. Install a compatible MindIE-SD build, select another "
+                "diffusion_kv_cache_dtype, or set it to 'auto'. Automatic precision "
+                "fallback is not performed."
+            ) from exc
+        kwargs = dict(precision=method, layout=layout, scale=self.softmax_scale)
+        # Apply Omni's deterministic Q/K rotation before every Dense quantized path.
+        if method in ("fp8", "mxfp8", "mxfp4"):
+            from vllm_omni.platforms.npu.quant.kv_quant_npu import get_quant_attention_rotation
 
-        layout = self.qkv_layout or "BNSD"
-        # Models pass (B, S, H, D); NPU fused op expects (B, N, S, D).
-        out = fp8_rotate_quant_fa(
-            query.transpose(1, 2),
-            key.transpose(1, 2),
-            value.transpose(1, 2),
-            layout=layout,
-            softmax_scale=self.softmax_scale,
-        )
-        return out.transpose(1, 2)
+            rotation = get_quant_attention_rotation(query.device, query.dtype, query.shape[-1])
+            kwargs.update(q_rot=rotation, k_rot=rotation)
+        logger.info_once("NPU attention uses MindIE-SD %s Runtime, layout=%s.", method, layout)
+        # Execution errors propagate. Never retry with another precision.
+        return runtime(query, key, value, **kwargs)
 
     def forward_fa_npu(
         self,

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -146,8 +146,10 @@ class RainFusionAttentionBackend(AttentionBackend):
             from mindiesd import sparse_attention
         except ImportError as exc:
             raise ValueError(_MISSING_MINDIESD) from exc
-        if not _supports_video_spans(sparse_attention):
-            raise ValueError(_INCOMPATIBLE_MINDIESD)
+        if not callable(sparse_attention):
+            raise ValueError(_MISSING_MINDIESD)
+        # Single-video Wan only needs the legacy geometry. Check video_spans
+        # when a forward actually carries multiple videos.
 
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
@@ -174,6 +176,8 @@ class RainFusionAttentionImpl(AttentionImpl):
     irregular video tail internally, retaining it outside the sparse blocks.
     """
 
+    _supported_kv_cache_dtypes = {"npu": {"fp8", "mxfp4"}}
+
     def __init__(
         self,
         num_heads: int,
@@ -222,6 +226,8 @@ class RainFusionAttentionImpl(AttentionImpl):
     def _validate_parallel_config(self) -> None:
         config = get_current_diffusion_config_or_none()
         parallel_config = getattr(config, "parallel_config", None)
+        if getattr(parallel_config, "allgather_degree", 1) > 1:
+            raise ValueError("RAINFUSION_ATTN requires full Q/K/V sequences; AllGather-KV SP is unsupported.")
         ring_degree = getattr(parallel_config, "ring_degree", 1)
         if ring_degree > 1:
             # Ring gives each rank a slice of the sequence, so block selection
@@ -262,7 +268,9 @@ class RainFusionAttentionImpl(AttentionImpl):
         plan = self._resolve_plan(attn_metadata)
         if plan is None:
             return self.dense_fallback.forward_npu(query, key, value, attn_metadata)
-        return self._forward_sparse_npu(query, key, value, plan)
+        if query.ndim != 4 or key.shape != query.shape or value.shape != query.shape or plan.used_len > query.shape[1]:
+            raise ValueError("RainFusion video geometry must fit identical full-sequence BSND Q/K/V tensors.")
+        return self._forward_sparse_npu(query, key, value, plan, attn_metadata)
 
     def _resolve_plan(self, attn_metadata: AttentionMetadata | None) -> RainFusionPlan | None:
         """Return the rf_v2 geometry, or None when this forward must stay dense."""
@@ -271,18 +279,17 @@ class RainFusionAttentionImpl(AttentionImpl):
             return None
         if self.layer_idx is not None and self.layer_idx in rf.skip_layers:
             return None
+        step_idx = total_steps = None
         if is_forward_context_available():
             step_idx = get_forward_context().denoise_step_idx
             total_steps = get_forward_context().total_denoise_steps
-            if step_idx is not None and step_idx < rf.start_step:
+        if (rf.start_step and step_idx is None) or (rf.end_step and (step_idx is None or total_steps is None)):
+            return None
+        if step_idx is not None:
+            if step_idx < rf.start_step:
                 return None
             # Tail fallback: keep the last ``end_step`` denoise steps dense.
-            if (
-                rf.end_step > 0
-                and step_idx is not None
-                and total_steps is not None
-                and step_idx >= total_steps - rf.end_step
-            ):
+            if rf.end_step > 0 and total_steps is not None and step_idx >= total_steps - rf.end_step:
                 return None
         if self.qkv_layout is None:
             # The sparse path reads the sequence off dim 1, which the tensors alone
@@ -308,7 +315,16 @@ class RainFusionAttentionImpl(AttentionImpl):
                 "model must publish AttentionMetadata.video_layout for the sequence to be sparsified."
             )
             return None
-        max_seqlen_q = attn_metadata.extra.get("max_seqlen_q")
+
+        # The sparse Runtime constructs its own block mask. It cannot preserve
+        # an arbitrary caller-supplied mask or piecewise visibility constraints.
+        mask = attn_metadata.attn_mask
+        padding_only = layout.used_len is not None
+        if attn_metadata.full_attn_spans is not None or (
+            mask is not None and not (padding_only and mask.ndim == 2 and mask.dtype == torch.bool)
+        ):
+            return None
+        max_seqlen_q = attn_metadata.extra.get("max_seqlen_q", layout.used_len)
         if max_seqlen_q is None:
             logger.warning_once(
                 "RAINFUSION_ATTN staying dense: attention metadata is missing max_seqlen_q, so the "
@@ -459,19 +475,39 @@ class RainFusionAttentionImpl(AttentionImpl):
         key: torch.Tensor,
         value: torch.Tensor,
         plan: RainFusionPlan,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         try:
             from mindiesd import sparse_attention
         except ImportError:
             raise ImportError(_MISSING_MINDIESD)
-        if self.rainfusion.precision != "bf16" and not _mindiesd_supports_precision():
-            raise RuntimeError(
-                f"block_sparse.precision={self.rainfusion.precision!r} requires MindIE-SD "
-                "with sparse_attention(precision=...) support; the installed mindiesd "
-                "silently ignores it and would run the BF16 path. Install a compatible "
-                "MindIE-SD release or use precision='bf16'."
+        if plan.video_spans is not None and not _supports_video_spans(sparse_attention):
+            raise ValueError(_INCOMPATIBLE_MINDIESD)
+        extra = attn_metadata.extra if attn_metadata else {}
+        requested = extra.get("kv_cache_dtype", self.rainfusion.precision)
+        reason = None
+        if requested in ("float", "bf16"):
+            precision = "bf16"
+        else:
+            precision = requested
+            if requested not in ("fp8", "mxfp4", "mix"):
+                reason = f"sparse {requested} is not supported"
+            elif not _mindiesd_supports_precision():
+                reason = "MindIE-SD sparse_attention must explicitly support precision"
+            elif plan.video_spans is not None:
+                reason = "quantized multi-video RainFusion is not supported"
+            elif requested in ("fp8", "mxfp4") and query.shape[0] != 1:
+                reason = "quantized Wan BSA requires batch size 1"
+            elif requested in ("fp8", "mxfp4") and (query.shape[-1] < 64 or query.shape[-1] & (query.shape[-1] - 1)):
+                reason = "BSA Hadamard rotations require a power-of-two head dimension of at least 64"
+        if reason is not None:
+            raise ValueError(
+                f"RainFusion precision {requested!r} is unavailable: {reason}. "
+                "Set diffusion_kv_cache_dtype or block_sparse.precision to a method supported by the installed "
+                "MindIE-SD, or disable attention quantization. Automatic precision fallback is not performed."
             )
 
+        logger.info_once("RainFusion uses MindIE-SD sparse attention, precision=%s.", precision)
         used = plan.used_len
         q, k, v = (tensor[:, :used] for tensor in (query, key, value))
         # Ulysses has already gathered the full sequence onto this rank and split
@@ -480,10 +516,10 @@ class RainFusionAttentionImpl(AttentionImpl):
             "scale": self.softmax_scale,
             "head_num": query.shape[-2],
             "input_layout": _INPUT_LAYOUT,
-            "inner_precise": _INNER_PRECISE,
+            "inner_precise": 4 if precision != "bf16" else _INNER_PRECISE,
             "block_size": _BLOCK_SIZE,
             "sparsity": self.rainfusion.sparsity,
-            "precision": self.rainfusion.precision,
+            "precision": precision,
         }
         if plan.video_spans is not None:
             out = sparse_attention(
@@ -497,7 +533,9 @@ class RainFusionAttentionImpl(AttentionImpl):
         else:
             assert plan.prefix_len is not None and plan.latent_shape is not None
             common_kwargs.update(
-                sparse_type="rf_v2",
+                # rf_v2 ignores precision on non-A5 devices. Select the explicit
+                # quantized high-level path so mask/rearrangement remain in MindIE.
+                sparse_type="rf_v3" if precision != "bf16" else "rf_v2",
                 txt_len=plan.prefix_len,
                 latent_shape_q=plan.latent_shape,
                 latent_shape_k=plan.latent_shape,

vllm_omni/diffusion/attention/layer.py

@@ -321,7 +321,7 @@ class Attention(nn.Module):
             dtype = None
         parallel_config = getattr(config, "parallel_config", None)
         ring_degree = getattr(parallel_config, "ring_degree", 1)
-        if dtype:
+        if dtype and dtype != "float":
             if ring_degree > 1:
                 raise ValueError(
                     "KV quantization is not compatible with ring attention "
@@ -338,13 +338,15 @@ class Attention(nn.Module):
         self._kv_cache_dtype = dtype
         self._kv_cache_skip_steps = getattr(config, "diffusion_kv_cache_skip_step_indices", None)
         self._kv_cache_skip_layers = getattr(config, "diffusion_kv_cache_skip_layer_indices", None)
+        if self._kv_cache_skip_layers and self.layer_idx is None and not self._disable_kv_quant:
+            raise ValueError("Attention quantization skip_layers requires a parseable transformer block index.")
 
     def _should_apply_kv_cache_quant(self) -> bool:
         skip_steps = self._kv_cache_skip_steps
         skip_layers = self._kv_cache_skip_layers
         if skip_steps is not None:
             step_idx = get_forward_context().denoise_step_idx if is_forward_context_available() else None
-            if step_idx is not None and step_idx in skip_steps:
+            if skip_steps and (step_idx is None or step_idx in skip_steps):
                 return False
         if skip_layers is not None:
             if self.layer_idx is not None and self.layer_idx in skip_layers:
@@ -352,18 +354,21 @@ class Attention(nn.Module):
         return True
 
     def _with_kv_cache_dtype(self, attn_metadata: AttentionMetadata | None) -> AttentionMetadata | None:
-        kv_cache_dtype = self._kv_cache_dtype
-        if kv_cache_dtype is None or self._disable_kv_quant or not self._should_apply_kv_cache_quant():
-            if attn_metadata is None or "kv_cache_dtype" not in attn_metadata.extra:
-                return attn_metadata
-            extra = dict(attn_metadata.extra)
-            extra.pop("kv_cache_dtype", None)
-            return replace(attn_metadata, extra=extra)
-
+        disabled = self._disable_kv_quant or not self._should_apply_kv_cache_quant()
+        dtype = self._kv_cache_dtype
+        if dtype in (None, "float"):
+            dtype = None
+        elif disabled:
+            dtype = "float"
+        if dtype is None and (attn_metadata is None or "kv_cache_dtype" not in attn_metadata.extra):
+            return attn_metadata
+        extra = dict(attn_metadata.extra) if attn_metadata is not None else {}
+        # Recompute per forward so shared metadata cannot retain another step's policy.
+        extra.pop("kv_cache_dtype", None)
+        if dtype is not None:
+            extra["kv_cache_dtype"] = dtype
         if attn_metadata is None:
-            return AttentionMetadata(extra={"kv_cache_dtype": kv_cache_dtype})
-        extra = dict(attn_metadata.extra)
-        extra["kv_cache_dtype"] = kv_cache_dtype
+            return AttentionMetadata(extra=extra) if extra else None
         return replace(attn_metadata, extra=extra)
 
     def forward(

vllm_omni/diffusion/data.py

@@ -1028,11 +1028,8 @@ class OmniDiffusionConfig:
     # has already resolved to vLLM's ModelOpt FP8 linear method.
     force_cutlass_fp8: bool = False
 
-    # Diffusion attention KV cache dtype (not vLLM's --kv-cache-dtype for AR models).
-    # None = native dtype (no quantization).
-    # "fp8" = dynamic FP8 (float8_e4m3fn) quantization per forward pass.
-    # On Hopper+FA3: native FP8 attention (memory + compute savings).
-    # On other backends: no benefit, backends skip quantization.
+    # Runtime diffusion attention method (not vLLM's --kv-cache-dtype for AR models).
+    # None/"auto" keeps native dtype; other values are validated by the selected backend.
     diffusion_kv_cache_dtype: str | None = None
     # Optional skip selectors for KV-cache quantization. Format: "0-9,20,25-30".
     # Listed steps/layers skip quantization; others keep quantized execution.

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -29,7 +29,7 @@ from vllm.model_executor.models.utils import (
 )
 from vllm.sequence import IntermediateTensors
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_pipeline_parallel_world_size,
@@ -417,6 +417,7 @@ class WanSelfAttention(nn.Module):
             softmax_scale=1.0 / (head_dim**0.5),
             causal=False,
             role="self",
+            qkv_layout="BSND",
             prefix=prefix,
         )
 
@@ -609,8 +610,8 @@ class WanCrossAttention(nn.Module):
             qkv_layout="BSND",
             prefix=prefix,
             skip_sequence_parallel=True,
-            # Wan2.2 cross-attn operates on short text-encoder sequences; per-block
-            # FP8 quant offers no perf win and degrades quality. Opt out until a
+            # Wan2.2 cross-attn operates on short text-encoder sequences; runtime
+            # Q/K/V quantization offers no perf win and degrades quality. Opt out until a
             # dedicated quant backend handles this case.
             disable_kv_quant=True,
         )
@@ -764,7 +765,15 @@ class WanTransformerBlock(nn.Module):
             self_attn_extra["vsa_dit_seq_shape"] = vsa_dit_seq_shape
         if preserve_vsa_all_blocks:
             self_attn_extra["preserve_vsa_all_blocks"] = True
-        self_attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask, extra=self_attn_extra)
+        video_layout = None
+        if vsa_dit_seq_shape is not None:
+            grid = tuple(int(dim) for dim in vsa_dit_seq_shape)
+            # Publish geometry independently of backend selection. Do not add
+            # partial cu_seqlens/max_seqlen metadata to dense CUDA attention.
+            video_layout = VideoTokenLayout(prefix_len=0, latent_grid=grid, used_len=math.prod(grid))
+        self_attn_metadata = AttentionMetadata(
+            attn_mask=hidden_states_mask, extra=self_attn_extra, video_layout=video_layout
+        )
         attn_output = self.attn1(norm_hidden_states, rotary_emb, self_attn_metadata)
         hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states)
 

vllm_omni/entrypoints/cli/serve.py

@@ -941,7 +941,8 @@ class OmniServeCommand(CLISubcommand):
             "--diffusion-kv-cache-dtype",
             type=str,
             default=None,
-            help="Diffusion attention KV cache dtype (e.g. fp8). Separate from vLLM --kv-cache-dtype.",
+            help="Diffusion Q/K/V precision: fp8, mxfp8, mxfp4, or float (NPU). "
+            "Separate from vLLM --kv-cache-dtype. Use --diffusion-attention-config for per-role fallback.",
         )
         omni_config_group.add_argument(
             "--diffusion-kv-cache-skip-steps",

vllm_omni/platforms/npu/quant/kv_quant_npu.py

@@ -1,127 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""FP8 quantization utilities for diffusion attention tensors.
-
-Provides per-tensor dynamic quantization of Q/K/V tensors to
-float8_e4m3fn format. Designed for diffusion models where Q/K/V are
-computed fresh each forward pass (no persistent KV cache).
-"""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Seeded rotation matrices for MindIE-SD quantized attention."""
 
 from __future__ import annotations
 
+import functools
 import math
-import threading
-from functools import lru_cache
 
 import torch
 
-# Hadamard rotation matrix for QuaRot-style preprocessing
-# keyed by (device, dtype, head_dim) to avoid matmul dtype mismatch.
-_ROT_MATRIXS: dict[tuple[torch.device, torch.dtype, int], torch.Tensor] = {}
-_ROT_MATRIX_LOCK = threading.Lock()
 
-_FP8_KV_LABELS = frozenset({"fp8"})
-
-
-def is_quantized_kv_cache(kv_cache_dtype: str | None) -> bool:
-    """True if config requests FP8-style KV / QKV quantization for the NPU FA path."""
-    return kv_cache_dtype in _FP8_KV_LABELS
-
-
-@lru_cache(maxsize=1)
-def _load_quant_ops():
-    try:
-        import torch_npu
-        from mindiesd.layers.quant.block_quant import fa_block_quant_preprocess
-        from msmodelslim.processor.quarot.common.quarot_utils import QuaRotMode, create_rot
-    except ImportError as e:
-        raise ImportError(
-            "fp8_rotate_quant_fa requires torch_npu, MindIE-SD (mindiesd), and MSModelSlim. "
-            "See https://gitcode.com/Ascend/MindIE-SD and https://gitcode.com/Ascend/msmodelslim"
-        ) from e
-    return torch_npu, fa_block_quant_preprocess, QuaRotMode, create_rot
-
-
-def _get_rot_matrix(
-    device: torch.device,
-    dtype: torch.dtype,
-    head_dim: int,
-    qua_rot_mode,
-    create_rot,
-) -> torch.Tensor:
-    key = (device, dtype, head_dim)
-    with _ROT_MATRIX_LOCK:
-        rot = _ROT_MATRIXS.get(key)
-        if rot is None:
-            rot = create_rot(qua_rot_mode.HADAMARD, head_dim, seed=425500).to(device=device, dtype=dtype)
-            _ROT_MATRIXS[key] = rot
-    return rot
-
-
-def fp8_rotate_quant_fa(
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    *,
-    layout: str = "BNSD",
-    softmax_scale: float | None = None,
-) -> torch.Tensor:
-    """Run NPU fused attention with dynamic FP8 Q/K/V and optional QuaRot preprocess.
-
-    Args:
-        query: Query tensor in ``layout`` order (default BNSD: batch, heads, seq, dim).
-        key: Key tensor in ``layout`` order (default BNSD: batch, heads, seq, dim).
-        value: Value tensor in ``layout`` order (default BNSD: batch, heads, seq, dim).
-        layout: ``BNSD`` or ``BSND`` for ``npu_fused_infer_attention_score_v2``.
-        softmax_scale: If None, uses ``1 / sqrt(head_dim)``.
-
-    Returns:
-        Attention output in the same layout as inputs.
-    """
-    torch_npu, fa_block_quant_preprocess, qua_rot_mode, create_rot = _load_quant_ops()
-
-    out_dtype = query.dtype
-    device = query.device
-
-    if layout == "BNSD":
-        _, n, s, d = query.shape
-    elif layout == "BSND":
-        _, s, n, d = query.shape
-    else:
-        raise ValueError(f"fp8_rotate_quant_fa: unsupported layout {layout!r}, expected BNSD or BSND")
-
-    rot = _get_rot_matrix(device, query.dtype, d, qua_rot_mode, create_rot)
-    q_f = torch.matmul(query, rot)
-    k_f = torch.matmul(key, rot)
-
-    q, q_scale = fa_block_quant_preprocess(q_f, block_size=128, dst_type=torch_npu.float8_e4m3fn, layout=layout)
-    k, k_scale = fa_block_quant_preprocess(k_f, block_size=256, dst_type=torch_npu.float8_e4m3fn, layout=layout)
-    v, v_scale = fa_block_quant_preprocess(value, block_size=256, dst_type=torch_npu.float8_e4m3fn, layout=layout)
-
-    scale = softmax_scale if softmax_scale is not None else 1.0 / math.sqrt(d)
-
-    out = torch_npu.npu_fused_infer_attention_score_v2(
-        q,
-        k,
-        v,
-        input_layout=layout,
-        num_query_heads=n,
-        softmax_scale=scale,
-        pre_tokens=2147483647,  # INT32_MAX: no left-context truncation.
-        next_tokens=2147483647,  # INT32_MAX: no right-context truncation.
-        query_quant_mode=7,  # NPU mode id for block FP8 dequant path.
-        key_quant_mode=7,  # Same quant mode as query branch.
-        value_quant_mode=7,  # Same quant mode as key/query branches.
-        dequant_scale_query=q_scale,
-        dequant_scale_key=k_scale,
-        dequant_scale_value=v_scale,
-        out_dtype=out_dtype,
-    )[0]
-
-    if out.shape[2] != s:
-        if layout == "BNSD":
-            out = out[:, :, :s, :]
-        elif layout == "BSND":
-            out = out[:, :s, :, :]
-
-    return out
+@functools.lru_cache(maxsize=128)
+def get_quant_attention_rotation(device: torch.device, dtype: torch.dtype, head_dim: int) -> torch.Tensor:
+    """Preserve Omni's fixed Walsh-Hadamard rotation without changing global RNG."""
+    if head_dim <= 0 or head_dim & (head_dim - 1):
+        raise ValueError("Generated attention rotations require a power-of-two head dimension.")
+    generator = torch.Generator(device="cpu").manual_seed(425500)
+    signs = torch.randint(0, 2, (head_dim,), generator=generator, dtype=torch.int64)
+    signs = signs.to(dtype=torch.float32).mul_(2).sub_(1)
+    hadamard = torch.ones(1, 1)
+    while hadamard.shape[0] < head_dim:
+        hadamard = torch.cat((torch.cat((hadamard, hadamard), dim=1), torch.cat((hadamard, -hadamard), dim=1)), dim=0)
+    rotation = signs[:, None] * hadamard / math.sqrt(head_dim)
+    return rotation.to(device=device, dtype=dtype).contiguous()