[AutoRound] Add offline quantized MXFP4 model support (#5544)

Signed-off-by: jl9876 <jia.li@intel.com>
Co-authored-by: Chendi.Xue <chendi.xue@intel.com>

tests/diffusion/quantization/test_wan_autoround_mxfp4.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Wan AutoRound-format MXFP4 configuration and layer-mapping tests."""
+
+from unittest.mock import patch
+
+import pytest
+from torch import nn
+from vllm.model_executor.layers.linear import LinearBase
+from vllm.model_executor.model_loader.utils import configure_quant_config
+
+from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import (
+    WanTransformer3DModel,
+)
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.inc_config import OmniINCConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _autoround_mxfp4_config() -> dict:
+    return {
+        "bits": 4,
+        "act_bits": 4,
+        "data_type": "mx_fp",
+        "act_data_type": "mx_fp",
+        "group_size": 32,
+        "act_group_size": 32,
+        "sym": True,
+        "act_sym": True,
+        "act_dynamic": True,
+        "enable_quanted_input": False,
+        "block_name_to_quantize": "blocks",
+        "quant_method": "auto-round",
+        "packing_format": "auto_round:llm_compressor",
+        "autoround_version": "0.14.2",
+    }
+
+
+def test_build_wan_autoround_mxfp4_config():
+    config = build_quant_config(_autoround_mxfp4_config())
+
+    assert isinstance(config, OmniINCConfig)
+    assert config.get_name() == "inc"
+    assert config.weight_bits == 4
+    assert config.data_type == "mx_fp"
+    assert config.group_size == 32
+    assert config.packing_format == "auto_round:llm_compressor"
+    assert config.block_name_to_quantize == ["blocks"]
+
+
+def test_wan_autoround_config_uses_runtime_layer_names():
+    config = build_quant_config(_autoround_mxfp4_config())
+    configure_quant_config(config, WanTransformer3DModel)
+
+    assert config.packed_modules_mapping is WanTransformer3DModel.packed_modules_mapping
+    assert config.block_name_to_quantize == ["blocks"]
+
+    layer = object.__new__(LinearBase)
+    nn.Module.__init__(layer)
+    for prefix in (
+        "blocks.0.ffn.net_0.proj",
+        "blocks.0.ffn.net_2",
+        "blocks.0.attn1.to_out",
+    ):
+        layer_config = config.config_parser.resolve(layer, prefix)
+        assert layer_config.quantized
+        assert layer_config.is_mxfp4
+
+
+def test_wan_autoround_scope_dispatches_only_blocks_to_mxfp4():
+    config = build_quant_config(_autoround_mxfp4_config())
+    configure_quant_config(config, WanTransformer3DModel)
+
+    layer = object.__new__(LinearBase)
+    nn.Module.__init__(layer)
+    layer_config = config.config_parser.resolve(layer, "blocks.0.ffn.net_0.proj")
+    outside_config = config.config_parser.resolve(layer, "condition_embedder.time_proj")
+
+    assert layer_config.quantized
+    assert layer_config.is_mxfp4
+    assert not outside_config.quantized
+
+    fake_method = object()
+    with patch(
+        "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_scheme.INCMxfp4Scheme.get_linear_method",
+        return_value=fake_method,
+    ):
+        method = config.get_quant_method(layer, "blocks.0.ffn.net_0.proj")
+    assert method is fake_method

tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E smoke test for Wan2.1 AutoRound MXFP4 inference."""
+
+import os
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+QUANTIZED_MODEL = os.environ.get(
+    "WAN21_AUTOROUND_MXFP4_MODEL",
+    "INCModel/Wan2.1-T2V-1.3B-Diffusers-MXFP4-AutoRound",
+)
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+
+@hardware_test(res={"xpu": "B60"})
+@pytest.mark.parametrize(
+    "omni_runner",
+    [(QUANTIZED_MODEL, None)],
+    indirect=True,
+)
+def test_wan21_autoround_mxfp4_generates_video(omni_runner, omni_runner_handler):
+    sampling_params = OmniDiffusionSamplingParams(
+        height=256,
+        width=256,
+        num_frames=5,
+        num_inference_steps=1,
+        guidance_scale=1.0,
+        generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+    )
+    response = omni_runner_handler.send_diffusion_request(
+        {
+            "prompt": "A red toy car driving on a table",
+            "sampling_params": sampling_params,
+        }
+    )
+
+    assert response.success, response.error_message
+    assert response.images is not None and len(response.images) == 1
+    frames = np.asarray(response.images[0])
+    assert frames.shape[1:4] == (5, 256, 256)
+    assert frames.std() > 0.01

vllm_omni/quantization/factory.py

@@ -470,10 +470,14 @@ def resolve_quant_config_from_disk(
         )
         return build_quant_config(qc_method, **qc_kwargs)
 
-    # AutoRound MXFP8 checkpoints use data_type="mx_fp" instead of
-    # is_checkpoint_*_serialized; rebuild so the offline path is selected.
+    # AutoRound MXFP checkpoints use data_type="mx_fp" instead of
+    # is_checkpoint_*_serialized; rebuild so the offline MXFP4/MXFP8 path is
+    # selected according to the checkpoint's bit width.
     if qc_kwargs.get("data_type") == "mx_fp":
-        logger.info("config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP8.")
+        logger.info(
+            "config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP%d.",
+            qc_kwargs.get("bits", getattr(quant_config, "weight_bits", 0)),
+        )
         return build_quant_config(qc_method, **qc_kwargs)
 
     if (