[Quant][AutoRound] Add W4A16 support for Qwen-Image (#3588)

Signed-off-by: yiliu30 <yi4.liu@intel.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E test for Qwen-Image AutoRound W4A16 quantized inference.
+
+Verifies that the W4A16 quantized Qwen-Image checkpoint loads end-to-end
+(quant_config propagated correctly) and produces a valid image.
+"""
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+QUANTIZED_MODEL = os.environ.get(
+    "QWEN_IMAGE_AUTOROUND_MODEL",
+    "INC4AI/Qwen-Image-AutoRound-W4A16",
+)
+
+HEIGHT = 512
+WIDTH = 512
+NUM_STEPS = 2
+
+
+def _sampling_params() -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        height=HEIGHT,
+        width=WIDTH,
+        num_inference_steps=NUM_STEPS,
+        true_cfg_scale=1.0,
+        generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+    )
+
+
+def _first_request_images(outputs) -> list:
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    req_out = first_output.request_output
+    assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+    return req_out.images
+
+
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"})
+def test_qwen_image_autoround_w4a16_load():
+    """Load the W4A16 quantized Qwen-Image model and run a minimal generation.
+
+    Verifies that quant_config is propagated to all transformer blocks
+    and produces a valid, non-blank image.
+    """
+    with OmniRunner(QUANTIZED_MODEL, enforce_eager=True) as runner:
+        outputs = runner.omni.generate(
+            "a cup of coffee on a table",
+            _sampling_params(),
+        )
+        images = _first_request_images(outputs)
+        assert len(images) >= 1, "Expected at least one generated image"
+        img = images[0]
+        assert img.width == WIDTH and img.height == HEIGHT
+        arr = np.array(img)
+        assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
+
+
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"})
+def test_qwen_image_autoround_w4a16_generate():
+    """Full generation: 512×512, 20 steps, CFG=5.0.
+
+    Validates that realistic quantized inference (multi-step with CFG
+    negative prompts and longer denoising) produces a valid non-blank image.
+    """
+    params = OmniDiffusionSamplingParams(
+        height=HEIGHT,
+        width=WIDTH,
+        num_inference_steps=20,
+        true_cfg_scale=5.0,
+        generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+    )
+    with OmniRunner(QUANTIZED_MODEL, enforce_eager=True) as runner:
+        outputs = runner.omni.generate(
+            "a cup of coffee on a wooden table, morning sunlight, photorealistic",
+            params,
+        )
+        images = _first_request_images(outputs)
+        assert len(images) >= 1, "Expected at least one generated image"
+        img = images[0]
+
+        # Save for manual inspection
+        output_path = Path(__file__).resolve().parents[4] / "ar_coffee_full.png"
+        img.save(str(output_path))
+
+        assert img.width == WIDTH and img.height == HEIGHT
+        arr = np.array(img)
+        assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"

vllm_omni/diffusion/models/flux/flux_transformer.py

@@ -27,6 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+from vllm_omni.quantization.component_config import safe_quant_config
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -44,23 +45,6 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
 logger = init_logger(__name__)
 
 
-def _safe_quant_config(quant_config: "QuantizationConfig | None") -> "QuantizationConfig | None":
-    """Return quant_config only if it is safe to propagate here, else None.
-
-    Dual-stream transformer_blocks, norm modulation layers, and norm_out are
-    kept at full precision for FP8 (see #2728). Offline quantization (e.g.
-    INC/AutoRound W4A16) needs the config propagated so packed weights load
-    correctly.
-    """
-    if quant_config is None:
-        return None
-    from vllm.model_executor.layers.quantization.inc import INCConfig
-
-    if isinstance(quant_config, INCConfig):
-        return quant_config
-    return None
-
-
 class ColumnParallelApproxGELU(nn.Module):
     def __init__(
         self,
@@ -401,7 +385,7 @@ class FluxSingleTransformerBlock(nn.Module):
         super().__init__()
         self.mlp_hidden_dim = int(dim * mlp_ratio)
 
-        self.norm = AdaLayerNormZeroSingle(dim, quant_config=_safe_quant_config(quant_config), prefix=f"{prefix}.norm")
+        self.norm = AdaLayerNormZeroSingle(dim, quant_config=safe_quant_config(quant_config), prefix=f"{prefix}.norm")
         self.proj_mlp = ReplicatedLinear(
             dim,
             self.mlp_hidden_dim,
@@ -596,7 +580,7 @@ class FluxTransformer2DModel(nn.Module):
                     dim=self.inner_dim,
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
-                    quant_config=_safe_quant_config(quant_config),
+                    quant_config=safe_quant_config(quant_config),
                     prefix=f"transformer_blocks.{i}",
                 )
                 for i in range(num_layers)
@@ -621,7 +605,7 @@ class FluxTransformer2DModel(nn.Module):
             self.inner_dim,
             elementwise_affine=False,
             eps=1e-6,
-            quant_config=_safe_quant_config(quant_config),
+            quant_config=safe_quant_config(quant_config),
             prefix="norm_out",
         )
         self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -332,7 +332,9 @@ class QwenImagePipeline(
         ).to(self.device)
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(
-            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+            od_config=od_config,
+            quant_config=od_config.quantization_config,
+            **transformer_kwargs,
         )
 
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -36,7 +36,6 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import calculate_
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
-from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.prompt_utils import (
@@ -257,6 +256,7 @@ class QwenImageEditPipeline(
         prefetch_subfolders(
             model,
             qwen_subfolders,
+            local_files_only=local_files_only,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
@@ -282,7 +282,11 @@ class QwenImageEditPipeline(
             local_files_only=local_files_only,
         ).to(self.device)
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
-        self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
+        self.transformer = QwenImageTransformer2DModel(
+            od_config=od_config,
+            quant_config=od_config.quantization_config,
+            **transformer_kwargs,
+        )
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
         self.processor = from_pretrained_with_prefetch(
             Qwen2VLProcessor.from_pretrained,
@@ -846,8 +850,10 @@ class QwenImageEditPipeline(
         if self.attention_kwargs is None:
             self._attention_kwargs = {}
 
-        txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds)
-        negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds)
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        negative_txt_seq_lens = (
+            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+        )
 
         latents = self.diffuse(
             prompt_embeds,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -250,7 +250,11 @@ class QwenImageEditPlusPipeline(
         ).to(self.device)
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
-        self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
+        self.transformer = QwenImageTransformer2DModel(
+            od_config=od_config,
+            quant_config=od_config.quantization_config,
+            **transformer_kwargs,
+        )
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
         self.processor = from_pretrained_with_prefetch(
             Qwen2VLProcessor.from_pretrained,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -275,7 +275,11 @@ class QwenImageLayeredPipeline(
         ]
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
-        self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
+        self.transformer = QwenImageTransformer2DModel(
+            od_config=od_config,
+            quant_config=od_config.quantization_config,
+            **transformer_kwargs,
+        )
 
         # Pipeline configuration & processing parameters
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -17,6 +17,7 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -25,6 +26,8 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.quantization.component_config import safe_quant_config
+
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
         QuantizationConfig,
@@ -48,8 +51,23 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 logger = init_logger(__name__)
 
 
-def _join_prefix(prefix: str, suffix: str) -> str:
-    return f"{prefix}.{suffix}" if prefix else suffix
+def _normalize_qwen_image_weight_name(name: str) -> str:
+    name = name.removeprefix("transformer.")
+    if ".to_out.0." in name:
+        name = name.replace(".to_out.0.", ".to_out.")
+    return name
+
+
+def _resolve_qwen_image_lookup_name(
+    name: str,
+    stacked_params_mapping: list[tuple[str, str, str]],
+) -> tuple[str, str | None]:
+    lookup_name = _normalize_qwen_image_weight_name(name)
+    for param_name, weight_name, shard_id in stacked_params_mapping:
+        if weight_name not in lookup_name or param_name in lookup_name:
+            continue
+        return lookup_name.replace(weight_name, param_name), shard_id
+    return lookup_name, None
 
 
 class ImageRopePrepare(nn.Module):
@@ -444,7 +462,7 @@ class ColumnParallelApproxGELU(nn.Module):
             gather_output=False,
             return_bias=False,
             quant_config=quant_config,
-            prefix=prefix,
+            prefix=f"{prefix}.proj" if prefix else "proj",
         )
         self.approximate = approximate
 
@@ -479,7 +497,7 @@ class FeedForward(nn.Module):
                 approximate="tanh",
                 bias=bias,
                 quant_config=quant_config,
-                prefix=_join_prefix(prefix, "net.0.proj"),
+                prefix=f"{prefix}.net.0",
             ),
             nn.Identity(),  # placeholder for weight loading
             RowParallelLinear(
@@ -488,7 +506,7 @@ class FeedForward(nn.Module):
                 input_is_parallel=True,
                 return_bias=False,
                 quant_config=quant_config,
-                prefix=_join_prefix(prefix, "net.2"),
+                prefix=f"{prefix}.net.2",
             ),
         ]
 
@@ -532,13 +550,13 @@ class QwenImageCrossAttention(nn.Module):
             head_size=self.head_dim,
             total_num_heads=num_heads,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "to_qkv"),
+            prefix=f"{prefix}.to_qkv",
         )
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -548,7 +566,7 @@ class QwenImageCrossAttention(nn.Module):
             head_size=head_dim,
             total_num_heads=num_heads,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "add_kv_proj"),
+            prefix=f"{prefix}.add_kv_proj",
         )
         self.add_query_num_heads = self.add_kv_proj.num_heads
         self.add_kv_num_heads = self.add_kv_proj.num_kv_heads
@@ -561,7 +579,7 @@ class QwenImageCrossAttention(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "to_add_out"),
+            prefix=f"{prefix}.to_add_out",
         )
 
         assert not pre_only
@@ -572,11 +590,11 @@ class QwenImageCrossAttention(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "to_out"),
+            prefix=f"{prefix}.to_out.0",
         )
 
-        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -713,22 +731,19 @@ class QwenImageTransformerBlock(nn.Module):
         self.dim = dim
         self.num_attention_heads = num_attention_heads
         self.attention_head_dim = attention_head_dim
+        txt_mod_quant_config = safe_quant_config(quant_config)
 
         # Image processing modules.
-        # Modulation linear is kept unquantized (quant_config=None) — it
-        # produces shift/scale/gate values that are precision-sensitive
-        # (see #2728). Use column TP with gather_output=True so weights are
-        # sharded while downstream modulation still receives full [B, 6 * dim].
+        # The re-quantized W4A16 checkpoint keeps img_mod.1 in full precision.
         self.img_mod = nn.Sequential(
             nn.SiLU(),
-            ColumnParallelLinear(
+            ReplicatedLinear(
                 dim,
                 6 * dim,
                 bias=True,
-                gather_output=True,
                 return_bias=False,
                 quant_config=None,
-                prefix=_join_prefix(prefix, "img_mod.1"),
+                prefix=f"{prefix}.img_mod.1",
             ),
         )
         self.img_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -739,27 +754,27 @@ class QwenImageTransformerBlock(nn.Module):
             context_pre_only=False,
             head_dim=attention_head_dim,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "attn"),
+            prefix=f"{prefix}.attn",
         )
         self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
         self.img_mlp = FeedForward(
             dim=dim,
             dim_out=dim,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "img_mlp"),
+            prefix=f"{prefix}.img_mlp",
         )
 
         # Text processing modules.
+        # AutoRound keeps txt_mod.1 quantized inside transformer_blocks.
         self.txt_mod = nn.Sequential(
             nn.SiLU(),
-            ColumnParallelLinear(
+            ReplicatedLinear(
                 dim,
                 6 * dim,
                 bias=True,
-                gather_output=True,
                 return_bias=False,
-                quant_config=None,
-                prefix=_join_prefix(prefix, "txt_mod.1"),
+                quant_config=txt_mod_quant_config,
+                prefix=f"{prefix}.txt_mod.1",
             ),
         )
         self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -769,7 +784,7 @@ class QwenImageTransformerBlock(nn.Module):
             dim=dim,
             dim_out=dim,
             quant_config=quant_config,
-            prefix=_join_prefix(prefix, "txt_mlp"),
+            prefix=f"{prefix}.txt_mlp",
         )
 
         self.zero_cond_t = zero_cond_t
@@ -979,6 +994,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
         self.out_channels = out_channels or in_channels
         self.inner_dim = num_attention_heads * attention_head_dim
         self.guidance_embeds = guidance_embeds
+        self.quant_config = quant_config
 
         if not use_layer3d_rope:
             self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
@@ -991,7 +1007,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).
@@ -1225,22 +1241,26 @@ class QwenImageTransformer2DModel(CachedTransformer):
 
         loaded_params: set[str] = set()
         for name, loaded_weight in weights:
-            original_name = name
-            lookup_name = name
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in original_name or param_name in original_name:
-                    continue
-                lookup_name = original_name.replace(weight_name, param_name)
-                param = params_dict[lookup_name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                if lookup_name not in params_dict and ".to_out.0." in lookup_name:
-                    lookup_name = lookup_name.replace(".to_out.0.", ".to_out.")
-                param = params_dict[lookup_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            original_name = name.removeprefix("transformer.")
+            lookup_name, shard_id = _resolve_qwen_image_lookup_name(
+                original_name,
+                stacked_params_mapping,
+            )
+
+            if lookup_name.endswith(".bias") and lookup_name not in params_dict:
+                continue
+
+            param = params_dict.get(lookup_name)
+            if param is None:
+                logger.debug("Skipping unexpected Qwen-Image transformer weight %s", original_name)
+                continue
+
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            if shard_id is None:
                 weight_loader(param, loaded_weight)
+            else:
+                weight_loader(param, loaded_weight, shard_id)
+
             loaded_params.add(original_name)
             loaded_params.add(lookup_name)
         return loaded_params

vllm_omni/quantization/component_config.py

@@ -48,6 +48,30 @@ def resolve_encoder_quant_config(
     return quant_config
 
 
+def safe_quant_config(
+    quant_config: QuantizationConfig | None,
+) -> QuantizationConfig | None:
+    """Return *quant_config* only if it is safe for norm/modulation layers.
+
+    Norm and modulation layers (LayerNorm, RMSNorm, AdaLayerNorm, img_mod,
+    txt_mod, etc.) produce precision-sensitive shift/scale/gate values and
+    should not receive FP8 quant configs (see #2728).  Pre-quantized methods
+    like INC/AutoRound W4A16 need the config propagated so packed weights
+    load correctly.
+
+    This is the inverse of :func:`resolve_encoder_quant_config`: that function
+    strips pre-quantized configs from encoders, while this one strips
+    *every config except* pre-quantized configs from norm/mod layers.
+    """
+    if quant_config is None:
+        return None
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    if isinstance(quant_config, INCConfig):
+        return quant_config
+    return None
+
+
 class ComponentQuantizationConfig(QuantizationConfig):
     """Routes quantization to different configs by layer prefix."""