[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_mxfp4_config.py

@@ -0,0 +1,539 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _patch_tp_state(monkeypatch):
+    """Patch TP rank/world_size so ModelWeightParameter can be instantiated on CPU
+    without an initialized distributed group.  Returns TP=1 rank=0 for all tests."""
+    monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP4Config
+# ---------------------------------------------------------------------------
+
+
+def test_mxfp4_config_get_name():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    assert DiffusionMXFP4Config.get_name() == "mxfp4"
+
+
+def test_mxfp4_config_from_config_defaults():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({})
+    assert cfg.is_checkpoint_mxfp4_serialized is False
+    assert cfg.ignored_layers == []
+
+
+def test_mxfp4_config_from_config_serialized():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({"is_checkpoint_mxfp4_serialized": True})
+    assert cfg.is_checkpoint_mxfp4_serialized is True
+
+
+def test_mxfp4_config_from_config_ignored_layers():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({"ignored_layers": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
+    """modules_to_not_convert must be accepted as an alias for ignored_layers."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = DiffusionMXFP4Config.from_config({"modules_to_not_convert": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+# ---------------------------------------------------------------------------
+# build_quant_config integration
+# ---------------------------------------------------------------------------
+
+
+def test_build_quant_config_mxfp4_string():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = build_quant_config("mxfp4")
+    assert isinstance(cfg, DiffusionMXFP4Config)
+    assert cfg.get_name() == "mxfp4"
+    assert cfg.is_checkpoint_mxfp4_serialized is False
+
+
+def test_build_quant_config_mxfp4_dict():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    cfg = build_quant_config({"method": "mxfp4", "is_checkpoint_mxfp4_serialized": True})
+    assert isinstance(cfg, DiffusionMXFP4Config)
+    assert cfg.is_checkpoint_mxfp4_serialized is True
+
+
+def test_build_quant_config_mxfp4_dualscale_string():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = build_quant_config("mxfp4_dualscale")
+    assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
+    assert cfg.is_checkpoint_serialized is False
+    assert cfg.num_bf16_fallback_layers == 5
+    assert cfg.ignored_layers == []
+
+
+def test_build_quant_config_mxfp4_dualscale_dict_offline():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = build_quant_config(
+        {
+            "method": "mxfp4_dualscale",
+            "is_checkpoint_serialized": True,
+            "ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"],
+        }
+    )
+    assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
+    assert cfg.is_checkpoint_serialized is True
+    assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"]
+
+
+def test_build_quant_config_mxfp4_dualscale_dict_online_custom_fallback():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = build_quant_config({"method": "mxfp4_dualscale", "num_bf16_fallback_layers": 10})
+    assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
+    assert cfg.num_bf16_fallback_layers == 10
+
+
+# ---------------------------------------------------------------------------
+# Block-index dispatch (_parse_block_idx)
+# ---------------------------------------------------------------------------
+
+
+def test_parse_block_idx_valid():
+    from vllm_omni.quantization.mxfp4_config import _parse_block_idx
+
+    assert _parse_block_idx("blocks.0.attn1.to_q") == 0
+    assert _parse_block_idx("blocks.5.ffn.net.0.proj") == 5
+    assert _parse_block_idx("blocks.40.norm1.weight") == 40
+
+
+def test_parse_block_idx_non_block_prefixes():
+    """Prefixes that do not start with 'blocks.N.' must return None."""
+    from vllm_omni.quantization.mxfp4_config import _parse_block_idx
+
+    assert _parse_block_idx("condition_embedder.time_embedder.linear_1") is None
+    assert _parse_block_idx("proj_out.weight") is None
+    assert _parse_block_idx("model.layers.0.self_attn.q_proj") is None
+    assert _parse_block_idx("scale_shift_table") is None
+
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_QUANTIZATION_METHODS
+# ---------------------------------------------------------------------------
+
+
+def test_supported_methods_include_mxfp4_variants():
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "mxfp4" in SUPPORTED_QUANTIZATION_METHODS
+    assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS
+    assert "mxfp4_dualscale" in SUPPORTED_QUANTIZATION_METHODS
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP4DualScaleMixedConfig — config roundtrips
+# ---------------------------------------------------------------------------
+
+
+def test_mixed_dualscale_config_get_name():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    assert DiffusionMXFP4DualScaleMixedConfig.get_name() == "mxfp4_dualscale"
+
+
+def test_mixed_dualscale_config_no_args_defaults():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig()
+    assert cfg.is_checkpoint_serialized is False
+    assert cfg.ignored_layers == []
+    assert cfg.num_bf16_fallback_layers == 5
+
+
+def test_mixed_dualscale_config_from_config_offline():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig.from_config(
+        {
+            "quant_method": "mxfp4_dualscale",
+            "is_checkpoint_serialized": True,
+            "ignored_layers": ["blocks.0.attn1.to_q", "proj_out"],
+        }
+    )
+    assert cfg.is_checkpoint_serialized is True
+    assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "proj_out"]
+    assert cfg.num_bf16_fallback_layers == 5  # default
+
+
+def test_mixed_dualscale_config_from_config_online_custom_fallback():
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig.from_config({"num_bf16_fallback_layers": 10})
+    assert cfg.is_checkpoint_serialized is False
+    assert cfg.num_bf16_fallback_layers == 10
+
+
+def test_mixed_dualscale_config_from_config_modules_to_not_convert_fallback():
+    """modules_to_not_convert must be accepted as an alias for ignored_layers."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig.from_config(
+        {"is_checkpoint_serialized": True, "modules_to_not_convert": ["proj_out"]}
+    )
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP4DualScaleMixedConfig — get_quant_method dispatch
+# ---------------------------------------------------------------------------
+
+
+def test_mixed_dualscale_offline_ignored_layer_returns_unquantized(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Offline: a prefix in ignored_layers must return UnquantizedLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=True,
+        ignored_layers=["blocks.0.attn1.to_q"],
+    )
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = cfg.get_quant_method(layer, "blocks.0.attn1.to_q")
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+def test_mixed_dualscale_offline_non_ignored_returns_mxfp4(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Offline: a prefix NOT in ignored_layers must return NPUMxfp4DualScaleLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=True,
+        ignored_layers=["blocks.0.attn1.to_q"],
+    )
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = cfg.get_quant_method(layer, "blocks.1.attn1.to_q")
+    assert isinstance(method, NPUMxfp4DualScaleLinearMethod)
+
+
+def test_mixed_dualscale_online_fallback_block_returns_unquantized(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: blocks < num_bf16_fallback_layers must return UnquantizedLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    assert isinstance(cfg.get_quant_method(layer, "blocks.0.attn1.to_q"), UnquantizedLinearMethod)
+    assert isinstance(cfg.get_quant_method(layer, "blocks.4.ffn.net.0.proj"), UnquantizedLinearMethod)
+
+
+def test_mixed_dualscale_online_quantized_block_returns_mxfp4(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: blocks >= num_bf16_fallback_layers must return NPUMxfp4DualScaleOnlineLinearMethod."""
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import (
+        DiffusionMXFP4DualScaleMixedConfig,
+        NPUMxfp4DualScaleOnlineLinearMethod,
+    )
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    assert isinstance(cfg.get_quant_method(layer, "blocks.5.attn1.to_q"), NPUMxfp4DualScaleOnlineLinearMethod)
+    assert isinstance(cfg.get_quant_method(layer, "blocks.40.ffn.net.0.proj"), NPUMxfp4DualScaleOnlineLinearMethod)
+
+
+def test_mixed_dualscale_online_non_block_prefix_returns_mxfp4(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: layers outside 'blocks.N.*' (condition_embedder etc.) always use MXFP4 online."""
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import (
+        DiffusionMXFP4DualScaleMixedConfig,
+        NPUMxfp4DualScaleOnlineLinearMethod,
+    )
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = cfg.get_quant_method(layer, "condition_embedder.time_embedder.linear_1")
+    assert isinstance(method, NPUMxfp4DualScaleOnlineLinearMethod)
+
+
+def test_mixed_dualscale_online_ignored_layers_override(
+    mocker,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Online: explicit ignored_layers must return UnquantizedLinearMethod regardless of block index.
+
+    A layer that is NOT in the leading-block range (block 10 >= num_bf16_fallback_layers=5)
+    but IS listed in ignored_layers must still fall back to BF16.  This lets power users
+    pin specific interleaved layers to BF16 during online quantization without needing an
+    offline checkpoint.
+    """
+    from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig(
+        is_checkpoint_serialized=False,
+        num_bf16_fallback_layers=5,
+        ignored_layers=["blocks.10.attn1.to_q"],
+    )
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    # block 10 is above the leading-block threshold but is in ignored_layers → BF16
+    assert isinstance(cfg.get_quant_method(layer, "blocks.10.attn1.to_q"), UnquantizedLinearMethod)
+
+
+def test_mixed_dualscale_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch):
+    """Non-LinearBase layers (norms, embeddings) must return None → no quantization."""
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
+
+    cfg = DiffusionMXFP4DualScaleMixedConfig()
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    norm_layer = torch.nn.LayerNorm(64)
+    assert cfg.get_quant_method(norm_layer, "blocks.0.norm1") is None
+
+
+# ---------------------------------------------------------------------------
+# TP=2 create_weights: parameter shapes and input_dim/output_dim
+#
+# Two scenarios mirror real Wan2.2 A14B linear layer types:
+#   Column-parallel (to_q, ffn.net_0): output is sharded (N/TP), input is full (K).
+#   Row-parallel   (to_out, ffn.net_2): input is sharded (K/TP), output is full (N).
+#
+# Tests verify:
+#   1. Registered parameter shapes are correct for each partition configuration.
+#   2. input_dim/output_dim attributes are set so RowParallelLinear.weight_loader
+#      can shard scale tensors correctly (the fix for the TP>1 shape-mismatch bug).
+#   3. Simulated loader slicing: slicing the full checkpoint tensor along the
+#      declared input_dim produces the exact shape stored in the parameter —
+#      proving the dim declaration is consistent with the allocation.
+# ---------------------------------------------------------------------------
+
+# K must be divisible by 32 (fine groups) and 512 (coarse groups).
+_TP2_K, _TP2_N, _TP2 = 1024, 512, 2
+
+
+class _FakeLayer(torch.nn.Module):
+    """Bare nn.Module that accepts register_parameter without a real weight_loader."""
+
+
+def _create_weights(method, *, input_size_per_partition, output_partition_sizes):
+    layer = _FakeLayer()
+    method.create_weights(
+        layer=layer,
+        input_size_per_partition=input_size_per_partition,
+        output_partition_sizes=output_partition_sizes,
+        input_size=_TP2_K,
+        output_size=_TP2_N,
+        params_dtype=torch.bfloat16,
+    )
+    return layer
+
+
+def _shard(tensor, param, rank, tp, dim_attr):
+    """Slice `tensor` along the dimension given by `param.<dim_attr>` for `rank`."""
+    dim = getattr(param, dim_attr)
+    if dim is None:
+        return tensor  # not sharded along this axis
+    shard_size = param.shape[dim]
+    slices = [slice(None)] * tensor.ndim
+    slices[dim] = slice(rank * shard_size, (rank + 1) * shard_size)
+    return tensor[tuple(slices)]
+
+
+# ── DualScale method ─────────────────────────────────────────────────────────
+
+
+def test_dualscale_column_parallel_tp2_shapes():
+    """Column-parallel TP=2: output halved, fine/coarse groups stay full, mul_scale full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
+
+    assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K)
+    assert layer.weight_scale.shape == (_TP2_N // _TP2, _TP2_K // 32)
+    assert layer.weight_dual_scale.shape == (_TP2_N // _TP2, _TP2_K // 512, 1)
+    assert layer.mul_scale.shape == (_TP2_K,)
+
+
+def test_dualscale_row_parallel_tp2_shapes():
+    """Row-parallel TP=2: input halved, fine/coarse groups halved, mul_scale halved."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
+    assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32)
+    assert layer.weight_dual_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 512, 1)
+    assert layer.mul_scale.shape == (_TP2_K // _TP2,)
+
+
+def test_dualscale_scale_parameter_input_dims():
+    """weight_scale/weight_dual_scale must have input_dim=1; mul_scale must have input_dim=0.
+
+    RowParallelLinear.weight_loader only shards a parameter when input_dim is set.
+    Without these, loading a full checkpoint tensor into a per-rank shape causes a
+    shape mismatch for TP>1 row-parallel layers (to_out, ffn.net_2).
+    """
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight_scale.input_dim == 1
+    assert layer.weight_scale.output_dim == 0
+    assert layer.weight_dual_scale.input_dim == 1
+    assert layer.weight_dual_scale.output_dim == 0
+    assert layer.mul_scale.input_dim == 0
+    assert layer.mul_scale.output_dim is None
+
+
+def test_dualscale_row_parallel_tp2_loader_simulation():
+    """Slicing full checkpoint tensors along input_dim must match row-parallel parameter shapes.
+
+    Simulates what RowParallelLinear.weight_loader does: for each scale parameter,
+    take the slice at rank*shard_size:(rank+1)*shard_size along input_dim.
+    The resulting shape must equal the per-rank parameter shape allocated by create_weights.
+    """
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    # Full checkpoint tensors (what the loader reads from disk).
+    ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
+    ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1)
+    ckpt_mul_scale = torch.zeros(_TP2_K)
+
+    for rank in range(_TP2):
+        assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape
+        assert (
+            _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "input_dim").shape
+            == layer.weight_dual_scale.shape
+        )
+        assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "input_dim").shape == layer.mul_scale.shape
+
+
+def test_dualscale_column_parallel_tp2_loader_simulation():
+    """Slicing full checkpoint tensors along output_dim must match column-parallel parameter shapes.
+
+    For column-parallel layers, the loader shards along output_dim (rows).
+    mul_scale has output_dim=None → not sharded (full tensor, same for all ranks).
+    """
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
+
+    method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
+
+    ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
+    ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1)
+    ckpt_mul_scale = torch.zeros(_TP2_K)
+
+    for rank in range(_TP2):
+        assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "output_dim").shape == layer.weight_scale.shape
+        assert (
+            _shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "output_dim").shape
+            == layer.weight_dual_scale.shape
+        )
+        # mul_scale: output_dim=None → no sharding → full tensor fits the column-parallel parameter
+        assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "output_dim").shape == layer.mul_scale.shape
+
+
+# ── Single-scale method ───────────────────────────────────────────────────────
+
+
+def test_single_scale_row_parallel_tp2_shapes():
+    """Row-parallel TP=2: input halved → weight_scale groups halved."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
+
+    method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
+    assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32)
+
+
+def test_single_scale_scale_parameter_input_dims():
+    """Single-scale weight_scale must have input_dim=1 for RowParallel TP sharding."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
+
+    method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight_scale.input_dim == 1
+    assert layer.weight_scale.output_dim == 0
+
+
+def test_single_scale_row_parallel_tp2_loader_simulation():
+    """Slicing full checkpoint weight_scale along input_dim matches row-parallel parameter shape."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
+
+    method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
+
+    for rank in range(_TP2):
+        assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape

tests/diffusion/quantization/test_mxfp4_key_remap.py

@@ -0,0 +1,399 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for merge_mxfp4_dualscale_checkpoint.py key-remapping helpers.
+
+These are pure-Python unit tests that exercise the transformation functions
+without loading any actual checkpoint files or requiring NPU hardware.
+"""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_MODEL_TYPES
+# ---------------------------------------------------------------------------
+
+
+def test_supported_model_types_includes_a14b():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_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_excludes_ti2v_5b():
+    """TI2V-5B is explicitly NOT supported under W4A4 quantization."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES
+
+    assert "Wan2.2-TI2V-5B" not in SUPPORTED_MODEL_TYPES
+
+
+# ---------------------------------------------------------------------------
+# _apply_rename_dict
+# ---------------------------------------------------------------------------
+
+
+def test_apply_rename_dict_self_attn_qkvo():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.0.self_attn.q.weight") == "blocks.0.attn1.to_q.weight"
+    assert _apply_rename_dict("blocks.0.self_attn.k.weight") == "blocks.0.attn1.to_k.weight"
+    assert _apply_rename_dict("blocks.0.self_attn.v.weight") == "blocks.0.attn1.to_v.weight"
+    assert _apply_rename_dict("blocks.0.self_attn.o.weight") == "blocks.0.attn1.to_out.0.weight"
+
+
+def test_apply_rename_dict_ffn():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.1.ffn.0.weight") == "blocks.1.ffn.net.0.proj.weight"
+    assert _apply_rename_dict("blocks.1.ffn.2.weight") == "blocks.1.ffn.net.2.weight"
+
+
+def test_apply_rename_dict_norm_order_swap():
+    """norm2↔norm3 swap: quant tool uses norm1/norm3/norm2 order,
+    Diffusers uses norm1/norm2/norm3."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.0.norm2.weight") == "blocks.0.norm3.weight"
+    assert _apply_rename_dict("blocks.0.norm3.weight") == "blocks.0.norm2.weight"
+
+
+def test_apply_rename_dict_cross_attn():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("blocks.0.cross_attn.q.weight") == "blocks.0.attn2.to_q.weight"
+    assert _apply_rename_dict("blocks.0.cross_attn.k.weight") == "blocks.0.attn2.to_k.weight"
+    assert _apply_rename_dict("blocks.0.cross_attn.v.weight") == "blocks.0.attn2.to_v.weight"
+    assert _apply_rename_dict("blocks.0.cross_attn.o.weight") == "blocks.0.attn2.to_out.0.weight"
+
+
+def test_apply_rename_dict_head_and_modulation():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
+
+    assert _apply_rename_dict("head.head.weight") == "proj_out.weight"
+
+
+# ---------------------------------------------------------------------------
+# _strip_mxfp4_wrapper
+# ---------------------------------------------------------------------------
+
+
+def test_strip_mxfp4_wrapper_linear_weight():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight") == "blocks.0.attn1.to_q.weight"
+
+
+def test_strip_mxfp4_wrapper_linear_weight_scale():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_scale") == "blocks.0.attn1.to_q.weight_scale"
+
+
+def test_strip_mxfp4_wrapper_linear_weight_dual_scale():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert (
+        _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_dual_scale") == "blocks.0.attn1.to_q.weight_dual_scale"
+    )
+
+
+def test_strip_mxfp4_wrapper_linear_bias():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.bias") == "blocks.0.attn1.to_q.bias"
+
+
+def test_strip_mxfp4_wrapper_div_mul_scale():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.div.mul_scale") == "blocks.0.attn1.to_q.mul_scale"
+
+
+def test_strip_mxfp4_wrapper_noop_for_plain_weight():
+    """MXFP8 / FLOAT tensors have no wrapper — must be returned unchanged."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
+
+    assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.weight") == "blocks.0.attn1.to_q.weight"
+    assert _strip_mxfp4_wrapper("blocks.0.norm_q.weight") == "blocks.0.norm_q.weight"
+    assert _strip_mxfp4_wrapper("condition_embedder.time_embedder.linear_1.weight") == (
+        "condition_embedder.time_embedder.linear_1.weight"
+    )
+
+
+# ---------------------------------------------------------------------------
+# _classify_blocks
+# ---------------------------------------------------------------------------
+
+
+def test_classify_blocks_mixed_mxfp8_and_mxfp4():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
+
+    quant_meta = {
+        "blocks.0.attn1.to_q.weight": "W8A8_MXFP8",
+        "blocks.1.attn1.to_q.weight": "W8A8_MXFP8",
+        "blocks.2.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
+        "blocks.3.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
+        "condition_embedder.time_embedder.linear_1.weight": "FLOAT",
+    }
+    block_types = _classify_blocks(quant_meta)
+    assert block_types[0] == "mxfp8"
+    assert block_types[1] == "mxfp8"
+    assert block_types[2] == "mxfp4_dualscale"
+    assert block_types[3] == "mxfp4_dualscale"
+    # Non-block key must not produce an entry
+    assert None not in block_types
+
+
+def test_classify_blocks_float_entries_skipped():
+    """FLOAT-typed tensors must not contribute to block classification."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
+
+    quant_meta = {
+        "blocks.0.bias": "FLOAT",
+        "blocks.1.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
+    }
+    block_types = _classify_blocks(quant_meta)
+    assert 0 not in block_types
+    assert block_types[1] == "mxfp4_dualscale"
+
+
+def test_classify_blocks_empty():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
+
+    assert _classify_blocks({}) == {}
+
+
+# ---------------------------------------------------------------------------
+# _is_mxfp4_tensor
+# ---------------------------------------------------------------------------
+
+
+def test_is_mxfp4_tensor_quantized_weight():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight", "W4A4_MXFP4_DUALSCALE") is True
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_scale", "W4A4_MXFP4_DUALSCALE") is True
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_dual_scale", "W4A4_MXFP4_DUALSCALE") is True
+
+
+def test_is_mxfp4_tensor_companion_bias_and_mul_scale():
+    """Companion tensors (bias, mul_scale) are FLOAT but belong to MXFP4 layers
+    and must be included in the merged checkpoint."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    # bias inside .linear. wrapper → must be included
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.bias", "FLOAT") is True
+    # mul_scale inside .div. wrapper → must be included
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.div.mul_scale", "FLOAT") is True
+
+
+def test_is_mxfp4_tensor_bf16_fallback_weight():
+    """BF16 fallback linear layers have plain .weight keys (no wrapper) → must NOT be included."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    assert _is_mxfp4_tensor("blocks.0.attn1.to_q.weight", "FLOAT") is False
+    assert _is_mxfp4_tensor("condition_embedder.time_embedder.linear_1.weight", "FLOAT") is False
+
+
+def test_is_mxfp4_tensor_norm_weight():
+    """Norm layers are always FLOAT with plain .weight → not MXFP4 related."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
+
+    assert _is_mxfp4_tensor("blocks.0.norm1.weight", "FLOAT") is False
+
+
+# ---------------------------------------------------------------------------
+# _collect_ignored_layers
+# ---------------------------------------------------------------------------
+
+
+def test_collect_ignored_layers_basic():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
+
+    # block 0: all three of Q/K/V are BF16 → fused as to_qkv in ignored_layers
+    # block 1: to_qkv is MXFP4 → not in ignored_layers
+    merged = {
+        "blocks.0.attn1.to_q.weight": torch.zeros(4, 4),
+        "blocks.0.attn1.to_k.weight": torch.zeros(4, 4),
+        "blocks.0.attn1.to_v.weight": torch.zeros(4, 4),
+        "blocks.1.attn1.to_q.weight": torch.zeros(4, 4),
+        "blocks.1.attn1.to_q.weight_scale": torch.zeros(4, 1),
+    }
+    mxfp4_prefixes = {"blocks.1.attn1.to_q"}
+
+    ignored = _collect_ignored_layers(merged, mxfp4_prefixes)
+    assert "blocks.0.attn1.to_qkv" in ignored  # fused vllm-omni name
+    assert "blocks.0.attn1.to_q" not in ignored  # diffusers name must not appear
+    assert "blocks.0.attn1.to_k" not in ignored
+    assert "blocks.1.attn1.to_qkv" not in ignored  # MXFP4, not ignored
+
+
+def test_collect_ignored_layers_empty_mxfp4():
+    """When no layers are MXFP4 (all BF16), all .weight prefixes become ignored_layers."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
+
+    # Use non-attn layers to avoid the Q/K/V completeness requirement.
+    merged = {
+        "blocks.0.ffn.net.0.proj.weight": torch.zeros(4, 4),
+        "proj_out.weight": torch.zeros(4, 4),
+    }
+    ignored = _collect_ignored_layers(merged, set())
+    assert sorted(ignored) == ["blocks.0.ffn.net_0.proj", "proj_out"]
+
+
+def test_collect_ignored_layers_returns_sorted():
+    """ignored_layers must be sorted for deterministic config.json output."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
+
+    # Use FFN layers (no Q/K/V fusion concern) to isolate the sort guarantee.
+    merged = {
+        "blocks.2.ffn.net.0.proj.weight": torch.zeros(1),
+        "blocks.0.ffn.net.0.proj.weight": torch.zeros(1),
+        "blocks.1.ffn.net.0.proj.weight": torch.zeros(1),
+    }
+    ignored = _collect_ignored_layers(merged, set())
+    assert ignored == sorted(ignored)
+
+
+# ---------------------------------------------------------------------------
+# _build_quant_config (new mxfp4_dualscale format)
+# ---------------------------------------------------------------------------
+
+
+def test_build_quant_config_new_format():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config
+
+    config = _build_quant_config(["blocks.0.attn1.to_q", "proj_out"])
+    assert config["quant_method"] == "mxfp4_dualscale"
+    assert config["is_checkpoint_serialized"] is True
+    assert config["ignored_layers"] == ["blocks.0.attn1.to_q", "proj_out"]
+    assert "num_mxfp8_blocks" not in config
+
+
+def test_build_quant_config_empty_ignored_layers():
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config
+
+    config = _build_quant_config([])
+    assert config["ignored_layers"] == []
+
+
+# ---------------------------------------------------------------------------
+# _diffusers_to_vllm_ignored
+# ---------------------------------------------------------------------------
+
+
+def test_remap_self_attn_qkv_fusion():
+    """All three of attn1.to_q/k/v in BF16 → fused attn1.to_qkv."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_k", "blocks.0.attn1.to_q", "blocks.0.attn1.to_v"])
+    assert result == ["blocks.0.attn1.to_qkv"]
+
+
+def test_remap_self_attn_qkv_partial_raises():
+    """Partial Q/K/V in ignored_layers is invalid — must raise ValueError.
+
+    Self-attention Q/K/V are fused into a single to_qkv layer at runtime;
+    partial precision (some BF16, some MXFP4) cannot be expressed and must
+    be caught early in the merge script.
+    """
+    import pytest
+
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    with pytest.raises(ValueError, match="Partial BF16 fallback"):
+        _diffusers_to_vllm_ignored(["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"])
+
+
+def test_remap_cross_attn_kept_separate():
+    """Cross-attention (attn2) to_q/k/v must NOT be fused."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.attn2.to_k", "blocks.0.attn2.to_q", "blocks.0.attn2.to_v"])
+    assert "blocks.0.attn2.to_qkv" not in result
+    assert "blocks.0.attn2.to_q" in result
+    assert "blocks.0.attn2.to_k" in result
+    assert "blocks.0.attn2.to_v" in result
+
+
+def test_remap_ffn_dot_to_underscore():
+    """ffn.net.0.proj → ffn.net_0.proj; ffn.net.2 → ffn.net_2."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.ffn.net.0.proj", "blocks.0.ffn.net.2"])
+    assert "blocks.0.ffn.net_0.proj" in result
+    assert "blocks.0.ffn.net_2" in result
+    assert "blocks.0.ffn.net.0.proj" not in result
+    assert "blocks.0.ffn.net.2" not in result
+
+
+def test_remap_to_out_strips_index():
+    """attn1.to_out.0 and attn2.to_out.0 must have the .0 suffix removed."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_out.0", "blocks.0.attn2.to_out.0"])
+    assert "blocks.0.attn1.to_out" in result
+    assert "blocks.0.attn2.to_out" in result
+    assert "blocks.0.attn1.to_out.0" not in result
+    assert "blocks.0.attn2.to_out.0" not in result
+
+
+def test_remap_noop_for_non_block_layers():
+    """Layers outside 'blocks.N.*' (condition_embedder, proj_out) pass through unchanged."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    inputs = [
+        "condition_embedder.time_embedder.linear_1",
+        "condition_embedder.text_embedder.linear_1",
+        "proj_out",
+    ]
+    result = _diffusers_to_vllm_ignored(inputs)
+    assert result == sorted(inputs)
+
+
+def test_remap_mixed_block_layers():
+    """Typical mixed block: QKV fused, to_out stripped, FFN renamed, cross-attn separate."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    inputs = [
+        "blocks.0.attn1.to_k",
+        "blocks.0.attn1.to_out.0",
+        "blocks.0.attn1.to_q",
+        "blocks.0.attn1.to_v",
+        "blocks.0.attn2.to_k",
+        "blocks.0.attn2.to_out.0",
+        "blocks.0.attn2.to_q",
+        "blocks.0.attn2.to_v",
+        "blocks.0.ffn.net.0.proj",
+        "blocks.0.ffn.net.2",
+    ]
+    result = _diffusers_to_vllm_ignored(inputs)
+    assert "blocks.0.attn1.to_qkv" in result
+    assert "blocks.0.attn1.to_out" in result
+    assert "blocks.0.attn2.to_q" in result
+    assert "blocks.0.attn2.to_k" in result
+    assert "blocks.0.attn2.to_v" in result
+    assert "blocks.0.attn2.to_out" in result
+    assert "blocks.0.ffn.net_0.proj" in result
+    assert "blocks.0.ffn.net_2" in result
+    # Old names must be gone
+    assert "blocks.0.attn1.to_q" not in result
+    assert "blocks.0.ffn.net.0.proj" not in result
+
+
+def test_remap_returns_sorted():
+    """Output of _diffusers_to_vllm_ignored must always be sorted."""
+    from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
+
+    inputs = [
+        "blocks.2.ffn.net.0.proj",
+        "blocks.0.attn1.to_q",
+        "blocks.0.attn1.to_k",
+        "blocks.0.attn1.to_v",
+        "blocks.1.attn2.to_out.0",
+    ]
+    result = _diffusers_to_vllm_ignored(inputs)
+    assert result == sorted(result)

tests/diffusion/quantization/test_mxfp8_config.py

@@ -0,0 +1,308 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for MXFP8 quantization config and linear method dispatch.
+
+Coverage:
+- DiffusionMXFP8Config.from_config roundtrips (CPU, no NPU required)
+- get_quant_method dispatch (mocked platform)
+- MXFPLinearMethodBase.apply() reshape skeleton (CPU)
+- Weight / scale shape-transform arithmetic from process_weights_after_loading (CPU)
+- build_quant_config integration
+- MXFP8_QUANT_CONFIG structure as the auto-detection contract
+"""
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available")
+
+
+# ---------------------------------------------------------------------------
+# DiffusionMXFP8Config — from_config roundtrips
+# ---------------------------------------------------------------------------
+
+
+def test_mxfp8_config_get_name():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    assert DiffusionMXFP8Config.get_name() == "mxfp8"
+
+
+def test_mxfp8_config_from_config_defaults():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({})
+    assert cfg.is_checkpoint_mxfp8_serialized is False
+    assert cfg.ignored_layers == []
+
+
+def test_mxfp8_config_from_config_serialized():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({"is_checkpoint_mxfp8_serialized": True})
+    assert cfg.is_checkpoint_mxfp8_serialized is True
+
+
+def test_mxfp8_config_from_config_ignored_layers():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({"ignored_layers": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+def test_mxfp8_config_from_config_modules_to_not_convert_fallback():
+    """modules_to_not_convert must be accepted as an alias for ignored_layers."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = DiffusionMXFP8Config.from_config({"modules_to_not_convert": ["proj_out"]})
+    assert cfg.ignored_layers == ["proj_out"]
+
+
+# ---------------------------------------------------------------------------
+# build_quant_config integration
+# ---------------------------------------------------------------------------
+
+
+def test_build_quant_config_mxfp8_string():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = build_quant_config("mxfp8")
+    assert isinstance(cfg, DiffusionMXFP8Config)
+    assert cfg.get_name() == "mxfp8"
+    assert cfg.is_checkpoint_mxfp8_serialized is False
+
+
+def test_build_quant_config_mxfp8_dict():
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    cfg = build_quant_config({"method": "mxfp8", "is_checkpoint_mxfp8_serialized": True})
+    assert isinstance(cfg, DiffusionMXFP8Config)
+    assert cfg.is_checkpoint_mxfp8_serialized is True
+
+
+def test_build_quant_config_mxfp8_config_json_format():
+    """Verify that the exact quantization_config injected by merge_mxfp8_checkpoint.py
+    is accepted by build_quant_config and selects the offline (serialized) path.
+
+    This is the critical auto-detection contract: TransformerConfig.from_dict()
+    reads quant_method + is_checkpoint_mxfp8_serialized to pick NPUMxfp8LinearMethod.
+    """
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
+
+    cfg = build_quant_config(MXFP8_QUANT_CONFIG)
+    assert isinstance(cfg, DiffusionMXFP8Config)
+    assert cfg.is_checkpoint_mxfp8_serialized is True
+
+
+def test_mxfp8_quant_config_structure():
+    """MXFP8_QUANT_CONFIG must contain exactly the keys that auto-detection reads."""
+    from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
+
+    assert MXFP8_QUANT_CONFIG.get("quant_method") == "mxfp8"
+    assert MXFP8_QUANT_CONFIG.get("is_checkpoint_mxfp8_serialized") is True
+
+
+# ---------------------------------------------------------------------------
+# get_quant_method dispatch
+# ---------------------------------------------------------------------------
+
+
+def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Offline (serialized) path on NPU must return NPUMxfp8LinearMethod."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod
+
+    config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
+    assert isinstance(method, NPUMxfp8LinearMethod)
+
+
+def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod
+
+    config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
+    assert isinstance(method, NPUMxfp8OnlineLinearMethod)
+
+
+def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Non-NPU platform must raise NotImplementedError (CUDA not yet supported)."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config()
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
+
+    with pytest.raises(NotImplementedError):
+        config.get_quant_method(layer, "blocks.0.attn1.to_q")
+
+
+def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """A prefix in ignored_layers must bypass quantization → UnquantizedLinearMethod."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config(ignored_layers=["proj_out"])
+    layer = mocker.Mock(spec=LinearBase)
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    method = config.get_quant_method(layer, "proj_out")
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+def test_get_quant_method_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch):
+    """Non-LinearBase layers (norms, embeddings) must get None → no quantization."""
+    from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
+
+    config = DiffusionMXFP8Config()
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
+
+    norm_layer = torch.nn.LayerNorm(64)
+    assert config.get_quant_method(norm_layer, "blocks.0.norm1") is None
+
+
+# ---------------------------------------------------------------------------
+# MXFPLinearMethodBase.apply() — reshape skeleton (CPU, no NPU)
+# ---------------------------------------------------------------------------
+
+
+def test_apply_reshape_skeleton():
+    """apply() must flatten batch dims → _apply_inner → restore original leading dims."""
+    from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase
+
+    OUT_FEATURES = 4
+
+    class _StubMethod(MXFPLinearMethodBase):
+        def create_weights(
+            self,
+            layer,
+            input_size_per_partition,
+            output_partition_sizes,
+            input_size,
+            output_size,
+            params_dtype,
+            **extra_weight_attrs,
+        ):
+            pass
+
+        def _quantize_activation(self, x):
+            return x, None
+
+        def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype):
+            return torch.zeros(x_q.shape[0], OUT_FEATURES, dtype=ori_dtype)
+
+    method = _StubMethod()
+    x = torch.randn(2, 3, 8)  # (batch=2, seq=3, K=8)
+    out = method.apply(None, x)
+    assert out.shape == (2, 3, OUT_FEATURES)
+
+
+def test_apply_reshape_with_bias():
+    """apply() must pass bias through to _apply_inner unchanged."""
+    from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase
+
+    received_bias = []
+
+    class _StubMethod(MXFPLinearMethodBase):
+        def create_weights(
+            self,
+            layer,
+            input_size_per_partition,
+            output_partition_sizes,
+            input_size,
+            output_size,
+            params_dtype,
+            **extra_weight_attrs,
+        ):
+            pass
+
+        def _quantize_activation(self, x):
+            return x, None
+
+        def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype):
+            received_bias.append(bias)
+            return torch.zeros(x_q.shape[0], 4, dtype=ori_dtype)
+
+    method = _StubMethod()
+    bias = torch.zeros(4)
+    method.apply(None, torch.randn(2, 8), bias=bias)
+    assert received_bias[0] is bias
+
+
+# ---------------------------------------------------------------------------
+# process_weights_after_loading shape arithmetic (pure torch, no NPU ops)
+#
+# These tests replicate the CPU-safe portions of process_weights_after_loading
+# to guard the key layout contract: (N,K) weight → (K,N) and (N,S) scale →
+# (S/2,N,2). They do NOT call NPU ops; they test only the math.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "n, k",
+    [(64, 128), (32, 64), (16, 96)],
+    ids=["64x128", "32x64", "16x96"],
+)
+def test_weight_transpose_contract(n: int, k: int):
+    """Weight must be transposed from (N, K) to (K, N) and be contiguous."""
+    w = torch.zeros(n, k, dtype=torch.uint8)
+    w = w.transpose(0, 1).contiguous()
+    assert w.shape == (k, n)
+    assert w.is_contiguous()
+
+
+@pytest.mark.parametrize(
+    "n, k_groups, expected_groups",
+    [
+        (64, 4, 2),  # even K_groups — no padding
+        (64, 3, 2),  # odd  K_groups — padded to 4
+        (32, 1, 1),  # odd  K_groups — padded to 2
+    ],
+    ids=["even", "odd-to-4", "odd-to-2"],
+)
+def test_weight_scale_reshape_contract(n: int, k_groups: int, expected_groups: int):
+    """Scale must be reshaped from (N, K_groups) to (K_groups_even//2, N, 2).
+
+    Odd K_groups must be padded to even before the reshape.
+    """
+    s = torch.zeros(n, k_groups, dtype=torch.uint8)
+    if k_groups % 2 == 1:
+        s = torch.cat([s, torch.zeros(n, 1, dtype=s.dtype)], dim=1)
+        k_groups += 1
+    s = s.reshape(n, k_groups // 2, 2).transpose(0, 1).contiguous()
+    assert s.shape == (expected_groups, n, 2)
+    assert s.is_contiguous()
+
+
+def test_num_groups_formula():
+    """K_groups formula: ceil(K / 32) — spot-check boundary values."""
+    assert (31 + 31) // 32 == 1  # K=31 → 1 group
+    assert (32 + 31) // 32 == 1  # K=32 → 1 group
+    assert (33 + 31) // 32 == 2  # K=33 → 2 groups
+    assert (128 + 31) // 32 == 4  # K=128 → 4 groups (even)
+    assert (96 + 31) // 32 == 3  # K=96  → 3 groups (odd → needs padding)
+
+
+# ---------------------------------------------------------------------------
+# SUPPORTED_QUANTIZATION_METHODS
+# ---------------------------------------------------------------------------
+
+
+def test_supported_methods_include_mxfp8():
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS

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)