From 0d5d1115d6f0b8be86b7b9f003705b6698540887 Mon Sep 17 00:00:00 2001 From: jl9876 Date: Tue, 25 Aug 2026 22:44:38 +0800 Subject: [PATCH] [AutoRound] Add offline quantized MXFP4 model support (#5544) Signed-off-by: jl9876 Co-authored-by: Chendi.Xue --- .../quantization/test_wan_autoround_mxfp4.py | 90 +++++++++++++++++++ .../test_wan21_autoround_mxfp4.py | 51 +++++++++++ vllm_omni/quantization/factory.py | 10 ++- 3 files changed, 148 insertions(+), 3 deletions(-) create mode 100644 tests/diffusion/quantization/test_wan_autoround_mxfp4.py create mode 100644 tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py diff --git a/tests/diffusion/quantization/test_wan_autoround_mxfp4.py b/tests/diffusion/quantization/test_wan_autoround_mxfp4.py new file mode 100644 index 000000000..89c317970 --- /dev/null +++ b/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 diff --git a/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py b/tests/e2e/offline_inference/test_wan21_autoround_mxfp4.py new file mode 100644 index 000000000..184f0f7d2 --- /dev/null +++ b/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 diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index b5d0a8ad7..245793871 100644 --- a/vllm_omni/quantization/factory.py +++ b/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 (