diff --git a/docs/user_guide/diffusion/attention_backends/rainfusion.md b/docs/user_guide/diffusion/attention_backends/rainfusion.md index 00f4e37a1..1d91cabfa 100644 --- a/docs/user_guide/diffusion/attention_backends/rainfusion.md +++ b/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. diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 8319077a2..7e7452bec 100644 --- a/docs/user_guide/quantization/overview.md +++ b/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 diff --git a/docs/user_guide/quantization/quantized_kvcache.md b/docs/user_guide/quantization/quantized_kvcache.md index 99a90e08a..30ba53593 100644 --- a/docs/user_guide/quantization/quantized_kvcache.md +++ b/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. diff --git a/tests/diffusion/attention/test_mindie_public_integration.py b/tests/diffusion/attention/test_mindie_public_integration.py new file mode 100644 index 000000000..c3df5351b --- /dev/null +++ b/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 diff --git a/tests/diffusion/attention/test_mindie_runtime.py b/tests/diffusion/attention/test_mindie_runtime.py new file mode 100644 index 000000000..bd9234a5f --- /dev/null +++ b/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" diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py index 5b5fcc405..6b364f1f1 100644 --- a/tests/diffusion/attention/test_rainfusion_plan.py +++ b/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) diff --git a/tests/diffusion/models/wan2_2/test_wan22_mindie_metadata.py b/tests/diffusion/models/wan2_2/test_wan22_mindie_metadata.py new file mode 100644 index 000000000..c316a207d --- /dev/null +++ b/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 diff --git a/tests/platforms/npu/quant/test_kv_quant_npu.py b/tests/platforms/npu/quant/test_kv_quant_npu.py index 0c828a6a3..611721809 100644 --- a/tests/platforms/npu/quant/test_kv_quant_npu.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index d1f09b51a..073fd3bd5 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py index 9bbde02fa..24a32bd31 100644 --- a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 2e0b2604d..0ad41e80f 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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( diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 6aa1a27d9..537148933 100644 --- a/vllm_omni/diffusion/data.py +++ b/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. diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 2808dcee7..91b482ebd 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/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) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 493f2533e..14e13cd4a 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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", diff --git a/vllm_omni/platforms/npu/quant/kv_quant_npu.py b/vllm_omni/platforms/npu/quant/kv_quant_npu.py index cb76aab51..76d090a7a 100644 --- a/vllm_omni/platforms/npu/quant/kv_quant_npu.py +++ b/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()