[NPU][Quant] Add W4A4 MXFP4 online & MXFP4 dual-scale online/offline quantization support for Wan2.2 T2V / I2V inference on Ascend NPU (#3578)

Signed-off-by: hyh_hh <huyinghong1@huawei.com>
Co-authored-by: hyh_hh <huyinghong1@huawei.com>

tests/diffusion/quantization/test_mxfp8_key_remap.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for merge_mxfp8_checkpoint.py key-remapping helpers and model metadata.
+
+These are pure-Python unit tests that exercise the transformation functions
+without loading any actual checkpoint files or requiring NPU hardware.
+
+Key contracts verified:
+- SUPPORTED_MODEL_TYPES includes Wan2.2-TI2V-5B (MXFP8 supports it; MXFP4 does not)
+- MXFP8_QUANT_CONFIG structure matches what TransformerConfig.from_dict() reads
+- _remap_keys correctly translates msModelSlim naming → Diffusers naming
+- _get_transformer_dirs routes cascade vs single-transformer models
+- _get_quant_subdir maps high/low noise subdirs for cascade models
+"""
+
+import pathlib
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_MODEL_TYPES and CASCADE_MODEL_TYPES
+# ---------------------------------------------------------------------------
+
+
+def test_supported_model_types_includes_all_wan22():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES
+
+    assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES
+    assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES
+
+
+def test_supported_model_types_includes_ti2v_5b():
+    """MXFP8 supports TI2V-5B (contrast: MXFP4 explicitly excludes it)."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES
+
+    assert "Wan2.2-TI2V-5B" in SUPPORTED_MODEL_TYPES
+
+
+def test_cascade_model_types_excludes_ti2v():
+    """TI2V-5B is a single-transformer model — must not be in CASCADE_MODEL_TYPES."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import CASCADE_MODEL_TYPES
+
+    assert "Wan2.2-TI2V-5B" not in CASCADE_MODEL_TYPES
+    assert "Wan2.2-T2V-A14B" in CASCADE_MODEL_TYPES
+    assert "Wan2.2-I2V-A14B" in CASCADE_MODEL_TYPES
+
+
+# ---------------------------------------------------------------------------
+# MXFP8_QUANT_CONFIG — auto-detection contract
+# ---------------------------------------------------------------------------
+
+
+def test_mxfp8_quant_config_has_required_keys():
+    """MXFP8_QUANT_CONFIG must carry exactly the keys that auto-detection reads:
+    quant_method (selects DiffusionMXFP8Config) and is_checkpoint_mxfp8_serialized
+    (selects NPUMxfp8LinearMethod over the online path)."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
+
+    assert MXFP8_QUANT_CONFIG["quant_method"] == "mxfp8"
+    assert MXFP8_QUANT_CONFIG["is_checkpoint_mxfp8_serialized"] is True
+
+
+# ---------------------------------------------------------------------------
+# _get_transformer_dirs
+# ---------------------------------------------------------------------------
+
+
+def test_get_transformer_dirs_cascade():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs
+
+    assert _get_transformer_dirs("Wan2.2-T2V-A14B") == ["transformer", "transformer_2"]
+    assert _get_transformer_dirs("Wan2.2-I2V-A14B") == ["transformer", "transformer_2"]
+
+
+def test_get_transformer_dirs_single():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs
+
+    assert _get_transformer_dirs("Wan2.2-TI2V-5B") == ["transformer"]
+
+
+# ---------------------------------------------------------------------------
+# _get_quant_subdir
+# ---------------------------------------------------------------------------
+
+
+def test_get_quant_subdir_cascade_high_noise():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
+
+    base = pathlib.Path("/quant")
+    result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer")
+    assert result == base / "high_noise_model"
+
+
+def test_get_quant_subdir_cascade_low_noise():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
+
+    base = pathlib.Path("/quant")
+    result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer_2")
+    assert result == base / "low_noise_model"
+
+
+def test_get_quant_subdir_non_cascade():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
+
+    base = pathlib.Path("/quant")
+    result = _get_quant_subdir("Wan2.2-TI2V-5B", base, "transformer")
+    assert result == base
+
+
+# ---------------------------------------------------------------------------
+# _remap_keys — msModelSlim naming → Diffusers naming
+# ---------------------------------------------------------------------------
+
+
+def test_remap_keys_self_attn_q():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"blocks.0.self_attn.q.weight": torch.zeros(1)}
+    meta = {"blocks.0.self_attn.q.weight": "W8A8_MXFP8"}
+    new_state, new_meta = _remap_keys(state, meta)
+    assert "blocks.0.attn1.to_q.weight" in new_state
+    assert new_meta.get("blocks.0.attn1.to_q.weight") == "W8A8_MXFP8"
+
+
+def test_remap_keys_self_attn_all_heads():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    pairs = {
+        "self_attn.q": "attn1.to_q",
+        "self_attn.k": "attn1.to_k",
+        "self_attn.v": "attn1.to_v",
+        "self_attn.o": "attn1.to_out.0",
+    }
+    for src_part, dst_part in pairs.items():
+        src_key = f"blocks.0.{src_part}.weight"
+        state = {src_key: torch.zeros(1)}
+        new_state, _ = _remap_keys(state, {})
+        expected = f"blocks.0.{dst_part}.weight"
+        assert expected in new_state, f"{src_key} → expected {expected}, got {list(new_state)}"
+
+
+def test_remap_keys_ffn():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {
+        "blocks.1.ffn.0.weight": torch.zeros(1),
+        "blocks.1.ffn.2.weight": torch.zeros(1),
+    }
+    new_state, _ = _remap_keys(state, {})
+    assert "blocks.1.ffn.net.0.proj.weight" in new_state
+    assert "blocks.1.ffn.net.2.weight" in new_state
+
+
+def test_remap_keys_norm_order_swap():
+    """norm2↔norm3 swap: msModelSlim uses norm1/norm3/norm2 order,
+    Diffusers uses norm1/norm2/norm3."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {
+        "blocks.0.norm2.weight": torch.zeros(1),
+        "blocks.0.norm3.weight": torch.zeros(1),
+    }
+    new_state, _ = _remap_keys(state, {})
+    # norm2 → norm3 and norm3 → norm2
+    assert "blocks.0.norm3.weight" in new_state
+    assert "blocks.0.norm2.weight" in new_state
+    # Both must be present (swap, not collapse)
+    assert len([k for k in new_state if "norm" in k and "norm_q" not in k and "norm_k" not in k]) == 2
+
+
+def test_remap_keys_head():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"head.head.weight": torch.zeros(1)}
+    new_state, _ = _remap_keys(state, {})
+    assert "proj_out.weight" in new_state
+
+
+def test_remap_keys_cross_attn():
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"blocks.0.cross_attn.q.weight": torch.zeros(1)}
+    new_state, _ = _remap_keys(state, {})
+    assert "blocks.0.attn2.to_q.weight" in new_state
+
+
+def test_remap_keys_meta_only_mapped_for_existing_state_keys():
+    """quant_meta entries are only emitted for keys present in state_dict."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    state = {"blocks.0.self_attn.q.weight": torch.zeros(1)}
+    # meta has an extra key not in state_dict
+    meta = {
+        "blocks.0.self_attn.q.weight": "W8A8_MXFP8",
+        "blocks.0.self_attn.q.weight_scale": "W8A8_MXFP8",
+    }
+    _, new_meta = _remap_keys(state, meta)
+    assert "blocks.0.attn1.to_q.weight" in new_meta
+    # weight_scale was in meta but not state_dict → must be absent
+    assert "blocks.0.attn1.to_q.weight_scale" not in new_meta
+
+
+def test_remap_keys_preserves_tensors():
+    """Tensor values must survive the key rename unchanged."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
+
+    t = torch.randn(4, 8)
+    state = {"blocks.0.self_attn.q.weight": t}
+    new_state, _ = _remap_keys(state, {})
+    assert torch.equal(new_state["blocks.0.attn1.to_q.weight"], t)