[Bugfix] Fix FLUX W4A16/AutoRound quant_config propagation (#3587)

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

tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py

@@ -26,10 +26,6 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
-pytestmark = pytest.mark.skip(
-    reason="Current Hugging Face account is not yet authorized to access this model.",
-)
-
 QUANTIZED_MODEL = "vllm-project-org/FLUX.1-dev-AutoRound-w4a16"
 BASELINE_MODEL = "black-forest-labs/FLUX.1-dev"
 
@@ -63,6 +59,27 @@ def _first_request_images(outputs) -> list:
     return req_out.images
 
 
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"})
+def test_flux_autoround_w4a16_load():
+    """Load the W4A16 quantized FLUX model and run a minimal generation.
+
+    Verifies that the quantized checkpoint loads end-to-end (quant_config
+    propagated to all transformer blocks) and produces a valid 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)"
+
+
 def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]:
     """Load a FLUX model, generate one image, return (images, peak_memory_mb)."""
     gc.collect()
@@ -87,6 +104,7 @@ def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]:
     return images, peak
 
 
+@pytest.mark.skip(reason="BF16 baseline model cannot be loaded correctly yet.")
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4"})
@@ -105,6 +123,7 @@ def test_flux_autoround_w4a16_generates_image(omni_runner: OmniRunner):
     assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
 
 
+@pytest.mark.skip(reason="BF16 baseline model cannot be loaded correctly yet.")
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4"})

vllm_omni/diffusion/models/flux/flux_transformer.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
+
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.layers.adalayernorm import (
@@ -41,6 +42,23 @@ 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,
@@ -381,9 +399,7 @@ class FluxSingleTransformerBlock(nn.Module):
         super().__init__()
         self.mlp_hidden_dim = int(dim * mlp_ratio)
 
-        # Modulation linear kept full precision; shift/scale/gate outputs
-        # are multiplied into the residual stream every block (see #2728).
-        self.norm = AdaLayerNormZeroSingle(dim, quant_config=None, 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,
@@ -565,16 +581,13 @@ class FluxTransformer2DModel(nn.Module):
         self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
         self.x_embedder = nn.Linear(in_channels, self.inner_dim)
 
-        # Dual-stream blocks kept full precision — FP8 on their joint
-        # attention path causes noise on FLUX (#2728). Single-stream
-        # blocks (38 vs 19) still get FP8 for memory savings.
         self.transformer_blocks = nn.ModuleList(
             [
                 FluxTransformerBlock(
                     dim=self.inner_dim,
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
-                    quant_config=None,
+                    quant_config=_safe_quant_config(quant_config),
                     prefix=f"transformer_blocks.{i}",
                 )
                 for i in range(num_layers)
@@ -594,13 +607,12 @@ class FluxTransformer2DModel(nn.Module):
             ]
         )
 
-        # Final modulation feeds proj_out; keep full precision (see #2728).
         self.norm_out = AdaLayerNormContinuous(
             self.inner_dim,
             self.inner_dim,
             elementwise_affine=False,
             eps=1e-6,
-            quant_config=None,
+            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)