[AutoRound] Add offline quantized W4A16 model support (#1777)

Signed-off-by: yiliu30 <yi4.liu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/user_guide/diffusion/quantization/autoround.md

@@ -0,0 +1,91 @@
+# AutoRound Quantization
+
+## Overview
+
+[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
+
+The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
+
+### Supported Schemes
+
+| Scheme | Bits | Status |
+|--------|------|--------|
+| W4A16 | 4 | ✅ Supported |
+| W8A16 | 8 | Planned |
+
+W4A16 is the first supported scheme. Additional schemes will be added in future releases.
+
+## Configuration
+
+1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=28),
+)
+outputs[0].save_images("output.png")
+```
+
+2. **CLI**: pass the quantized model path directly.
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
+  --prompt "A cat sitting on a windowsill" \
+  --num-inference-steps 28 \
+  --output outputs/flux_w4a16.png
+```
+
+No `--quantization` flag is needed — the quantization method is read from the checkpoint.
+
+## How It Works
+
+The checkpoint's `config.json` contains:
+
+```json
+{
+  "quantization_config": {
+    "quant_method": "auto-round",
+    "bits": 4,
+    "group_size": 128,
+    "sym": true,
+    "packing_format": "auto_round:auto_gptq",
+    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
+  }
+}
+```
+
+At load time:
+
+1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
+2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
+3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
+
+## Supported Models
+
+| Model | HF Checkpoint | Scheme | Group Size | Backend |
+|-------|--------------|--------|------------|---------|
+| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
+
+## Creating a Quantized Checkpoint
+
+Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
+
+```bash
+# W4A16 (4-bit weight, 16-bit activation)
+auto-round \
+    --model black-forest-labs/FLUX.1-dev \
+    --scheme W4A16 \
+    --batch_size 1 \
+    --disable_opt_rtn \
+    --dataset coco2014 \
+    --iters 0
+```
+
+The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.

docs/user_guide/diffusion/quantization/overview.md

@@ -11,6 +11,7 @@ vLLM-Omni provides a unified quantization framework that supports both diffusion
 | FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
 | Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
 | GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
+| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |
 
 ### Multi-stage Omni Models (Pre-quantized Checkpoints)
 
@@ -102,6 +103,9 @@ config = build_quant_config("fp8")
 # Dict with parameters
 config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
 
+# AutoRound / INC (auto-detected from checkpoint, or explicit)
+config = build_quant_config("auto-round", bits=4, group_size=128)
+
 # Per-component dict
 config = build_quant_config({
     "transformer": {"method": "fp8"},

tests/diffusion/layers/test_adalayernorm.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for shared AdaLayerNorm layers used by FLUX and other models."""
+
+import os
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Initialize the minimal distributed environment required by
+    ReplicatedLinear (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29501")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(autouse=True)
+def _force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
+
+    vLLM's dispatch_unquantized_gemm() selects the backend by platform
+    (e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
+    CPU test tensors crash with NotImplementedError on ROCm.  Monkeypatch
+    the dispatcher to always return the default (torch.nn.functional.linear)
+    implementation which works on any device."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def test_adalayernorm_import_from_shared_module():
+    """Verify imports work from the shared adalayernorm module."""
+    from vllm_omni.diffusion.layers.adalayernorm import (  # noqa: F401
+        AdaLayerNormContinuous,
+        AdaLayerNormZero,
+        AdaLayerNormZeroSingle,
+    )
+
+
+def test_adalayernorm_zero_forward_shape():
+    """AdaLayerNormZero produces correct output shapes (x, gate, shift, scale, gate)."""
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
+
+    dim = 64
+    batch = 2
+    seq_len = 4
+    norm = AdaLayerNormZero(dim)
+
+    x = torch.randn(batch, seq_len, dim)
+    emb = torch.randn(batch, dim)
+
+    out_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = norm(x, emb)
+
+    assert out_x.shape == (batch, seq_len, dim)
+    assert gate_msa.shape == (batch, dim)
+    assert shift_mlp.shape == (batch, dim)
+    assert scale_mlp.shape == (batch, dim)
+    assert gate_mlp.shape == (batch, dim)
+
+
+def test_adalayernorm_zero_single_forward_shape():
+    """AdaLayerNormZeroSingle produces (x, gate) with correct shapes."""
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
+
+    dim = 64
+    batch = 2
+    seq_len = 4
+    norm = AdaLayerNormZeroSingle(dim)
+
+    x = torch.randn(batch, seq_len, dim)
+    emb = torch.randn(batch, dim)
+
+    out_x, gate = norm(x, emb)
+
+    assert out_x.shape == (batch, seq_len, dim)
+    assert gate.shape == (batch, dim)
+
+
+def test_adalayernorm_continuous_forward_shape():
+    """AdaLayerNormContinuous produces correct output shape."""
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
+
+    dim = 64
+    cond_dim = 64
+    batch = 2
+    seq_len = 4
+    norm = AdaLayerNormContinuous(dim, cond_dim)
+
+    x = torch.randn(batch, seq_len, dim)
+    conditioning = torch.randn(batch, cond_dim)
+
+    out = norm(x, conditioning)
+
+    assert out.shape == (batch, seq_len, dim)
+
+
+def test_adalayernorm_zero_accepts_quant_config():
+    """Constructor accepts quant_config=None and prefix='test' without error."""
+    from vllm_omni.diffusion.layers.adalayernorm import (
+        AdaLayerNormContinuous,
+        AdaLayerNormZero,
+        AdaLayerNormZeroSingle,
+    )
+
+    # Should not raise with quant_config=None and prefix
+    AdaLayerNormZero(64, quant_config=None, prefix="test.norm1")
+    AdaLayerNormZeroSingle(64, quant_config=None, prefix="test.norm")
+    AdaLayerNormContinuous(64, 64, quant_config=None, prefix="test.norm_out")
+
+
+def test_adalayernorm_uses_replicated_linear():
+    """Verify .linear is a ReplicatedLinear instance (not nn.Linear)."""
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    from vllm_omni.diffusion.layers.adalayernorm import (
+        AdaLayerNormContinuous,
+        AdaLayerNormZero,
+        AdaLayerNormZeroSingle,
+    )
+
+    norm_zero = AdaLayerNormZero(64)
+    assert isinstance(norm_zero.linear, ReplicatedLinear)
+
+    norm_zero_single = AdaLayerNormZeroSingle(64)
+    assert isinstance(norm_zero_single.linear, ReplicatedLinear)
+
+    norm_continuous = AdaLayerNormContinuous(64, 64)
+    assert isinstance(norm_continuous.linear, ReplicatedLinear)
+
+
+# ── Numerical equivalence tests against diffusers originals ──
+
+
+def _copy_weights(src_linear, dst_replicated_linear):
+    """Copy weights from nn.Linear to ReplicatedLinear for comparison."""
+    dst_replicated_linear.weight.data.copy_(src_linear.weight.data)
+    if src_linear.bias is not None and dst_replicated_linear.bias is not None:
+        dst_replicated_linear.bias.data.copy_(src_linear.bias.data)
+
+
+def test_adalayernorm_zero_matches_diffusers():
+    """Verify AdaLayerNormZero produces identical output to diffusers original."""
+    from diffusers.models.normalization import (
+        AdaLayerNormZero as DiffusersAdaLayerNormZero,
+    )
+
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
+
+    dim = 64
+    torch.manual_seed(42)
+    ours = AdaLayerNormZero(dim)
+    ref = DiffusersAdaLayerNormZero(dim)
+
+    # Copy weights: nn.Linear -> ReplicatedLinear
+    _copy_weights(ref.linear, ours.linear)
+
+    x = torch.randn(2, 4, dim)
+    emb = torch.randn(2, dim)
+
+    out_ours = ours(x, emb)
+    out_ref = ref(x, emb=emb)
+
+    for o, r in zip(out_ours, out_ref):
+        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
+
+
+def test_adalayernorm_zero_single_matches_diffusers():
+    """Verify AdaLayerNormZeroSingle produces identical output to diffusers original."""
+    from diffusers.models.normalization import (
+        AdaLayerNormZeroSingle as DiffusersAdaLayerNormZeroSingle,
+    )
+
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
+
+    dim = 64
+    torch.manual_seed(42)
+    ours = AdaLayerNormZeroSingle(dim)
+    ref = DiffusersAdaLayerNormZeroSingle(dim)
+
+    _copy_weights(ref.linear, ours.linear)
+
+    x = torch.randn(2, 4, dim)
+    emb = torch.randn(2, dim)
+
+    out_ours = ours(x, emb)
+    out_ref = ref(x, emb=emb)
+
+    for o, r in zip(out_ours, out_ref):
+        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
+
+
+def test_adalayernorm_continuous_matches_diffusers():
+    """Verify AdaLayerNormContinuous produces identical output to diffusers original."""
+    from diffusers.models.normalization import (
+        AdaLayerNormContinuous as DiffusersAdaLayerNormContinuous,
+    )
+
+    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
+
+    dim = 64
+    cond_dim = 64
+    torch.manual_seed(42)
+    # Match constructor args: diffusers defaults elementwise_affine=True, eps=1e-5
+    ours = AdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
+    ref = DiffusersAdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
+
+    _copy_weights(ref.linear, ours.linear)
+
+    x = torch.randn(2, 4, dim)
+    cond = torch.randn(2, cond_dim)
+
+    out_ours = ours(x, cond)
+    out_ref = ref(x, cond)
+
+    torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5)

tests/diffusion/models/flux/test_flux_prefix_propagation.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests that FLUX transformer blocks correctly propagate `quant_config` and
+`prefix` through all sub-layers.
+
+The tests instantiate blocks with a known prefix and verify that all quantization-
+aware sub-layers (AdaLayerNorm, FeedForward, Attention projections) receive the
+prefix rooted at the block prefix.  This is critical for quantized weight loading
+to match checkpoint keys to the correct model parameters.
+"""
+
+import os
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+# Standard dimensions for a minimal FLUX block
+_DIM = 64
+_HEADS = 2
+_HEAD_DIM = _DIM // _HEADS
+
+
+@pytest.fixture(autouse=True)
+def _init_distributed():
+    """Initialize the minimal distributed environment required by
+    vLLM parallel linear layers (tensor-parallel group must exist)."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29502")
+    init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method="env://",
+    )
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+def _param_names(module) -> set[str]:
+    """Return the set of all parameter names in a module."""
+    return {name for name, _ in module.named_parameters()}
+
+
+def test_flux_transformer_block_passes_prefix():
+    """FluxTransformerBlock propagates prefix to norm1, norm1_context, attn, ff, ff_context."""
+    from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformerBlock
+
+    prefix = "transformer_blocks.0"
+    block = FluxTransformerBlock(
+        dim=_DIM,
+        num_attention_heads=_HEADS,
+        attention_head_dim=_HEAD_DIM,
+        quant_config=None,
+        prefix=prefix,
+    )
+
+    params = _param_names(block)
+
+    # norm1 and norm1_context (AdaLayerNormZero) should have linear weights
+    assert any(name.startswith("norm1.linear.") for name in params), (
+        f"norm1.linear.* not found in params: {sorted(params)}"
+    )
+    assert any(name.startswith("norm1_context.linear.") for name in params), (
+        f"norm1_context.linear.* not found in params: {sorted(params)}"
+    )
+
+    # attn should have QKV projections
+    assert any(name.startswith("attn.to_qkv.") for name in params), (
+        f"attn.to_qkv.* not found in params: {sorted(params)}"
+    )
+
+    # ff and ff_context should have net layers
+    assert any(name.startswith("ff.net.") for name in params), f"ff.net.* not found in params: {sorted(params)}"
+    assert any(name.startswith("ff_context.net.") for name in params), (
+        f"ff_context.net.* not found in params: {sorted(params)}"
+    )
+
+
+def test_flux_single_transformer_block_passes_prefix():
+    """FluxSingleTransformerBlock propagates prefix to norm, proj_mlp, attn."""
+    from vllm_omni.diffusion.models.flux.flux_transformer import FluxSingleTransformerBlock
+
+    prefix = "single_transformer_blocks.0"
+    block = FluxSingleTransformerBlock(
+        dim=_DIM,
+        num_attention_heads=_HEADS,
+        attention_head_dim=_HEAD_DIM,
+        quant_config=None,
+        prefix=prefix,
+    )
+
+    params = _param_names(block)
+
+    # norm (AdaLayerNormZeroSingle) should have linear weights
+    assert any(name.startswith("norm.linear.") for name in params), (
+        f"norm.linear.* not found in params: {sorted(params)}"
+    )
+
+    # proj_mlp (ReplicatedLinear) should have weight
+    assert any(name.startswith("proj_mlp.") for name in params), f"proj_mlp.* not found in params: {sorted(params)}"
+
+    # attn should have QKV projection
+    assert any(name.startswith("attn.to_qkv.") for name in params), (
+        f"attn.to_qkv.* not found in params: {sorted(params)}"
+    )
+
+
+def test_flux_feedforward_passes_prefix():
+    """FeedForward propagates prefix to net.0 (GELU proj) and net.2 (output proj)."""
+    from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
+
+    prefix = "transformer_blocks.0.ff"
+    ff = FeedForward(
+        dim=_DIM,
+        dim_out=_DIM,
+        quant_config=None,
+        prefix=prefix,
+    )
+
+    params = _param_names(ff)
+
+    # net.0 is ColumnParallelApproxGELU which wraps a ColumnParallelLinear
+    assert any("net.0" in name for name in params), f"net.0 not found in params: {sorted(params)}"
+
+    # net.2 is RowParallelLinear
+    assert any("net.2" in name for name in params), f"net.2 not found in params: {sorted(params)}"

tests/diffusion/quantization/test_inc_config.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for INC/AutoRound quantization via the unified framework."""
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+def test_build_quant_config_autoround():
+    """build_quant_config("auto-round", ...) should produce an INCConfig."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config(
+        "auto-round",
+        bits=4,
+        group_size=128,
+        sym=True,
+        packing_format="auto_round:auto_gptq",
+    )
+    assert config is not None
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+    assert config.group_size == 128
+
+
+def test_build_quant_config_inc():
+    """build_quant_config("inc", ...) should also produce an INCConfig."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("inc", bits=4, group_size=128)
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+
+
+def test_build_quant_config_autoround_dict():
+    """Dict-style config with method=auto-round should work."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config(
+        {
+            "method": "auto-round",
+            "bits": 4,
+            "group_size": 128,
+            "sym": True,
+            "packing_format": "auto_round:auto_gptq",
+        }
+    )
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+
+
+def test_build_quant_config_autoround_filters_metadata():
+    """Checkpoint metadata keys (autoround_version, batch_size, iters)
+    should be silently filtered out instead of causing TypeError."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config(
+        "auto-round",
+        bits=4,
+        group_size=128,
+        sym=True,
+        packing_format="auto_round:auto_gptq",
+        block_name_to_quantize="transformer_blocks,single_transformer_blocks",
+        autoround_version="0.12.0",  # metadata — must be filtered
+        batch_size=1,  # metadata — must be filtered
+        iters=0,  # metadata — must be filtered
+    )
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+    assert config.group_size == 128
+
+
+def test_build_quant_config_bits_to_weight_bits_mapping():
+    """The 'bits' key from checkpoints should be mapped to 'weight_bits'."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    # If weight_bits is already provided, bits should be ignored
+    config = build_quant_config("auto-round", weight_bits=4, group_size=128)
+    assert isinstance(config, INCConfig)
+    assert config.weight_bits == 4
+
+
+def test_autoround_in_supported_methods():
+    """auto-round and inc should appear in SUPPORTED_QUANTIZATION_METHODS."""
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "auto-round" in SUPPORTED_QUANTIZATION_METHODS
+    assert "inc" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_integration_autoround_via_omni_diffusion_config():
+    """OmniDiffusionConfig with auto-round quantization dict should resolve."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "method": "auto-round",
+            "bits": 4,
+            "group_size": 128,
+            "sym": True,
+        },
+    )
+    assert isinstance(config.quantization_config, INCConfig)
+    assert config.quantization_config.weight_bits == 4
+
+
+def test_integration_autodetect_from_transformer_config():
+    """When TransformerConfig has quant_config, OmniDiffusionConfig should
+    auto-detect it even without explicit quantization_config."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
+
+    tf_config = TransformerConfig.from_dict(
+        {
+            "quantization_config": {
+                "quant_method": "auto-round",
+                "bits": 4,
+                "group_size": 128,
+                "sym": True,
+                "packing_format": "auto_round:auto_gptq",
+                "autoround_version": "0.12.0",
+                "batch_size": 1,
+                "iters": 0,
+            }
+        }
+    )
+    assert tf_config.quant_method == "auto-round"
+    assert isinstance(tf_config.quant_config, INCConfig)
+
+    od_config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
+    assert isinstance(od_config.quantization_config, INCConfig)
+    assert od_config.quantization_config.weight_bits == 4

tests/e2e/offline_inference/test_flux_autoround_w4a16.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E tests for FLUX AutoRound W4A16 quantized inference.
+
+These tests require:
+  - A CUDA GPU
+  - The quantized model checkpoint (vllm-project-org/FLUX.1-dev-AutoRound-w4a16)
+"""
+
+import gc
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+
+from tests.utils import DeviceMemoryMonitor, hardware_test
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+
+QUANTIZED_MODEL = "vllm-project-org/FLUX.1-dev-AutoRound-w4a16"
+BASELINE_MODEL = "black-forest-labs/FLUX.1-dev"
+
+# Allow overriding via environment for local testing
+import os as _os
+
+QUANTIZED_MODEL = _os.environ.get("FLUX_AUTOROUND_MODEL", QUANTIZED_MODEL)
+BASELINE_MODEL = _os.environ.get("FLUX_BASELINE_MODEL", BASELINE_MODEL)
+
+# Small resolution to keep GPU memory & time manageable
+HEIGHT = 256
+WIDTH = 256
+NUM_STEPS = 2  # minimal for smoke-test
+
+
+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()
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    current_omni_platform.reset_peak_memory_stats()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+
+    m = Omni(model=model_name, enforce_eager=True, **extra_kwargs)
+
+    current_omni_platform.reset_peak_memory_stats()
+    outputs = m.generate(
+        "a photo of a cat sitting on a laptop keyboard",
+        OmniDiffusionSamplingParams(
+            height=HEIGHT,
+            width=WIDTH,
+            num_inference_steps=NUM_STEPS,
+            guidance_scale=0.0,
+            generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
+        ),
+    )
+
+    peak = monitor.peak_used_mb
+    monitor.stop()
+
+    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")
+    images = req_out.images
+
+    del m
+    gc.collect()
+    current_omni_platform.empty_cache()
+
+    return images, peak
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"})
+def test_flux_autoround_w4a16_generates_image():
+    """Load the W4A16 quantized FLUX model and verify it produces a valid image."""
+    images, _ = _generate_image(QUANTIZED_MODEL)
+
+    assert len(images) >= 1, "Expected at least one generated image"
+    img = images[0]
+    assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
+    assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
+
+    # Sanity: image should not be blank (all-zero)
+    import numpy as np
+
+    arr = np.array(img)
+    assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"})
+def test_flux_autoround_w4a16_memory_savings():
+    """Compare peak GPU memory of quantized vs FP16 baseline.
+
+    The W4A16 model should use meaningfully less memory than the
+    BF16/FP16 baseline since weights are 4-bit instead of 16-bit.
+    """
+    quant_images, quant_peak = _generate_image(QUANTIZED_MODEL)
+    cleanup_dist_env_and_memory()
+    _, baseline_peak = _generate_image(BASELINE_MODEL)
+
+    print(f"Quantized (W4A16) peak memory: {quant_peak:.0f} MB")
+    print(f"Baseline (BF16) peak memory:   {baseline_peak:.0f} MB")
+    print(f"Savings:                        {baseline_peak - quant_peak:.0f} MB")
+
+    # W4A16 weights are 4x smaller than BF16/FP16.  FLUX.1-dev transformer
+    # is ~12 GB in BF16, so we expect ~9 GB savings on weights alone.
+    # Use a conservative threshold to account for activations and overhead.
+    min_savings_mb = 2000
+    assert quant_peak + min_savings_mb < baseline_peak, (
+        f"Quantized model ({quant_peak:.0f} MB) should use at least "
+        f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)"
+    )

vllm_omni/diffusion/data.py

@@ -193,12 +193,24 @@ class TransformerConfig:
     """Container for raw transformer configuration dictionaries."""
 
     params: dict[str, Any] = field(default_factory=dict)
+    quant_method: str | None = None
+    quant_config: "QuantizationConfig | None" = None
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "TransformerConfig":
         if not isinstance(data, dict):
             raise TypeError(f"Expected transformer config dict, got {type(data)!r}")
-        return cls(params=dict(data))
+        params = dict(data)  # copy to avoid mutating caller's dict
+
+        quant_method: str | None = None
+        quant_config: QuantizationConfig | None = None
+        disk_qc = params.get("quantization_config")
+        if isinstance(disk_qc, dict) and "quant_method" in disk_qc:
+            quant_method = disk_qc["quant_method"]
+            kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
+            quant_config = build_quant_config(quant_method, **kwargs)
+
+        return cls(params=params, quant_method=quant_method, quant_config=quant_config)
 
     def to_dict(self) -> dict[str, Any]:
         return dict(self.params)
@@ -598,6 +610,17 @@ class OmniDiffusionConfig:
             # If it's neither dict nor DiffusionCacheConfig, convert to empty config
             self.cache_config = DiffusionCacheConfig()
 
+        # Auto-detect quantization from TransformerConfig if not explicitly set.
+        # This covers the case where tf_model_config is passed at construction
+        # time.  For late (post-construction) assignment, callers should use
+        # set_tf_model_config() which propagates quant_config automatically.
+        if self.quantization_config is None and self.tf_model_config.quant_config is not None:
+            self.quantization_config = self.tf_model_config.quant_config
+            logger.info(
+                "Auto-detected quantization '%s' from model config",
+                self.tf_model_config.quant_method,
+            )
+
         # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.
         if self.quantization_config is not None:
             if isinstance(self.quantization_config, QuantizationConfig):
@@ -617,6 +640,28 @@ class OmniDiffusionConfig:
         elif self.max_cpu_loras < 1:
             raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
 
+    def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
+        """Assign `tf_model_config` and propagate quantization if detected.
+
+        In the normal startup flow `OmniDiffusionConfig` is created
+        *before* the transformer `config.json` is loaded from disk, so
+        `__post_init__` sees an empty `TransformerConfig`.  Callers
+        that load the config later should use this method instead of bare
+        assignment so that an embedded `quant_config` is propagated to
+        `self.quantization_config` automatically.
+
+        Args:
+            tf_config: Transformer configuration, typically built via
+                `TransformerConfig.from_dict`.
+        """
+        self.tf_model_config = tf_config
+        if self.quantization_config is None and tf_config.quant_config is not None:
+            self.quantization_config = tf_config.quant_config
+            logger.info(
+                "Auto-detected quantization '%s' from model config",
+                tf_config.quant_method,
+            )
+
     def update_multimodal_support(self) -> None:
         self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"}
 

vllm_omni/diffusion/layers/adalayernorm.py

@@ -1,11 +1,16 @@
 from importlib.util import find_spec
+from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ReplicatedLinear
 
 from vllm_omni.diffusion.layers.custom_op import CustomOp
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
 logger = init_logger(__name__)
 
 _HAS_MINDIESD = find_spec("mindiesd") is not None
@@ -123,3 +128,102 @@ class AdaLayerNorm(CustomOp):
         shift_result, scale_result, gate_result = self.preprocess(mod_params, index)
 
         return self.layernorm(x) * (1 + scale_result) + shift_result, gate_result
+
+
+class AdaLayerNormZero(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.emb = None
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(
+            embedding_dim,
+            6 * embedding_dim,
+            bias=bias,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
+        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        emb = self.linear(self.silu(emb))
+        if isinstance(emb, tuple):
+            emb = emb[0]
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class AdaLayerNormZeroSingle(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(
+            embedding_dim,
+            3 * embedding_dim,
+            bias=bias,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
+        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        emb = self.linear(self.silu(emb))
+        if isinstance(emb, tuple):
+            emb = emb[0]
+        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa
+
+
+class AdaLayerNormContinuous(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        elementwise_affine: bool = False,
+        eps: float = 1e-6,
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(
+            conditioning_embedding_dim,
+            embedding_dim * 2,
+            bias=bias,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.linear",
+        )
+        self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
+
+    def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
+        emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
+        if isinstance(emb, tuple):
+            emb = emb[0]
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -7,7 +7,7 @@ import re
 import time
 from collections.abc import Generator, Iterable
 from pathlib import Path
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import torch
 from huggingface_hub import hf_hub_download
@@ -34,6 +34,9 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.registry import initialize_model
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
 logger = init_logger(__name__)
 
 
@@ -332,12 +335,60 @@ class DiffusersPipelineLoader:
             weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")}
             weights_not_loaded = weights_not_loaded - weights_scale_not_loaded
             if weights_not_loaded:
-                raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
+                self._check_unloaded_weights(weights_not_loaded)
             if weights_scale_not_loaded:
                 logger.warning(
                     f"Following weight_scale weights were not initialized from checkpoint: {weights_scale_not_loaded}"
                 )
 
+    @staticmethod
+    def _is_expected_quantized_weight(name: str) -> bool:
+        """Return True if *name* is a quantization-specific parameter.
+
+        Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra
+        parameters that have no counterpart in an unquantized checkpoint.
+        These are expected to be absent and should not trigger a load error.
+        """
+        # Weight suffixes that quantization methods register in the model but
+        # are not present in unquantized checkpoints.
+        _QUANTIZED_WEIGHT_SUFFIXES = (
+            # GPTQ / AWQ / AutoRound – g_idx is optional (not all checkpoints include it)
+            ".g_idx",
+            # FP8
+            ".weight_scale",
+            ".weight_scale_inv",
+            ".input_scale",
+            # GGUF
+            ".qweight_type",
+            # INT8  (weight_scale already covered above)
+        )
+        return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES)
+
+    def _check_unloaded_weights(
+        self,
+        weights_not_loaded: set[str],
+    ) -> None:
+        """Validate unloaded weights, tolerating expected quantization artifacts.
+
+        For quantized models, weights matching known quant-specific suffixes
+        are logged as a warning.  Any *other* missing weight raises
+        ``ValueError`` regardless of quantization.
+        """
+        od_config = getattr(self, "od_config", None)
+        if od_config is None or od_config.quantization_config is None:
+            raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
+
+        expected_missing = {w for w in weights_not_loaded if self._is_expected_quantized_weight(w)}
+        unexpected_missing = weights_not_loaded - expected_missing
+
+        if expected_missing:
+            logger.warning(
+                "Following weights were not initialized from checkpoint (expected for quantized models): %s",
+                expected_missing,
+            )
+        if unexpected_missing:
+            raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
+
     def _is_gguf_quantization(self, od_config: OmniDiffusionConfig) -> bool:
         quant_config = od_config.quantization_config
         if quant_config is None:

vllm_omni/diffusion/models/flux/flux_transformer.py

@@ -12,7 +12,6 @@ from diffusers.models.embeddings import (
     get_1d_rotary_pos_embed,
 )
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
 from diffusers.utils import is_torch_npu_available
 from vllm.distributed import get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather
 from vllm.logger import init_logger
@@ -32,6 +31,11 @@ if TYPE_CHECKING:
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.layers.adalayernorm import (
+    AdaLayerNormContinuous,
+    AdaLayerNormZero,
+    AdaLayerNormZeroSingle,
+)
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
 
 logger = init_logger(__name__)
@@ -46,6 +50,7 @@ class ColumnParallelApproxGELU(nn.Module):
         approximate: str,
         bias: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.proj = ColumnParallelLinear(
@@ -55,6 +60,7 @@ class ColumnParallelApproxGELU(nn.Module):
             gather_output=False,
             return_bias=False,
             quant_config=quant_config,
+            prefix=f"{prefix}.proj",
         )
         self.approximate = approximate
 
@@ -73,6 +79,7 @@ class FeedForward(nn.Module):
         inner_dim: int | None = None,
         bias: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
 
@@ -82,7 +89,9 @@ class FeedForward(nn.Module):
         dim_out = dim_out or dim
 
         layers: list[nn.Module] = [
-            ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config),
+            ColumnParallelApproxGELU(
+                dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config, prefix=f"{prefix}.net.0"
+            ),
             nn.Identity(),  # placeholder for weight loading
             RowParallelLinear(
                 inner_dim,
@@ -90,6 +99,7 @@ class FeedForward(nn.Module):
                 input_is_parallel=True,
                 return_bias=False,
                 quant_config=quant_config,
+                prefix=f"{prefix}.net.2",
             ),
         ]
 
@@ -117,6 +127,7 @@ class FluxAttention(torch.nn.Module):
         context_pre_only: bool | None = None,
         pre_only: bool = False,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
 
@@ -141,6 +152,7 @@ class FluxAttention(torch.nn.Module):
             total_num_heads=self.heads,
             bias=bias,
             quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv",
         )
 
         if not self.pre_only:
@@ -153,6 +165,7 @@ class FluxAttention(torch.nn.Module):
                         input_is_parallel=True,
                         return_bias=False,
                         quant_config=quant_config,
+                        prefix=f"{prefix}.to_out.0",
                     ),
                     nn.Dropout(dropout),
                 ]
@@ -168,8 +181,8 @@ class FluxAttention(torch.nn.Module):
                 total_num_heads=self.heads,
                 bias=added_proj_bias,
                 quant_config=quant_config,
+                prefix=f"{prefix}.add_kv_proj",
             )
-
             self.to_add_out = RowParallelLinear(
                 self.inner_dim,
                 query_dim,
@@ -177,6 +190,7 @@ class FluxAttention(torch.nn.Module):
                 input_is_parallel=True,
                 return_bias=False,
                 quant_config=quant_config,
+                prefix=f"{prefix}.to_add_out",
             )
 
         self.rope = RotaryEmbedding(is_neox_style=False)
@@ -272,11 +286,11 @@ class FluxTransformerBlock(nn.Module):
         qk_norm: str = "rms_norm",
         eps: float = 1e-6,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
-
-        self.norm1 = AdaLayerNormZero(dim)
-        self.norm1_context = AdaLayerNormZero(dim)
+        self.norm1 = AdaLayerNormZero(dim, quant_config=quant_config, prefix=f"{prefix}.norm1")
+        self.norm1_context = AdaLayerNormZero(dim, quant_config=quant_config, prefix=f"{prefix}.norm1_context")
 
         self.attn = FluxAttention(
             query_dim=dim,
@@ -288,13 +302,14 @@ class FluxTransformerBlock(nn.Module):
             bias=True,
             eps=eps,
             quant_config=quant_config,
+            prefix=f"{prefix}.attn",
         )
 
         self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config)
+        self.ff = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix=f"{prefix}.ff")
 
         self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.ff_context = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config)
+        self.ff_context = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix=f"{prefix}.ff_context")
 
     def forward(
         self,
@@ -361,17 +376,28 @@ class FluxSingleTransformerBlock(nn.Module):
         attention_head_dim: int,
         mlp_ratio: float = 4.0,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.mlp_hidden_dim = int(dim * mlp_ratio)
 
-        self.norm = AdaLayerNormZeroSingle(dim)
+        self.norm = AdaLayerNormZeroSingle(dim, quant_config=quant_config, prefix=f"{prefix}.norm")
         self.proj_mlp = ReplicatedLinear(
-            dim, self.mlp_hidden_dim, bias=True, return_bias=False, quant_config=quant_config
+            dim,
+            self.mlp_hidden_dim,
+            bias=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj_mlp",
         )
         self.act_mlp = nn.GELU(approximate="tanh")
         self.proj_out = ReplicatedLinear(
-            dim + self.mlp_hidden_dim, dim, bias=True, return_bias=False, quant_config=quant_config
+            dim + self.mlp_hidden_dim,
+            dim,
+            bias=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj_out",
         )
 
         self.attn = FluxAttention(
@@ -383,6 +409,7 @@ class FluxSingleTransformerBlock(nn.Module):
             eps=1e-6,
             pre_only=True,
             quant_config=quant_config,
+            prefix=f"{prefix}.attn",
         )
 
     def forward(
@@ -542,8 +569,9 @@ class FluxTransformer2DModel(nn.Module):
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
                     quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 
@@ -554,12 +582,20 @@ class FluxTransformer2DModel(nn.Module):
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
                     quant_config=quant_config,
+                    prefix=f"single_transformer_blocks.{i}",
                 )
-                for _ in range(num_single_layers)
+                for i in range(num_single_layers)
             ]
         )
 
-        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim,
+            self.inner_dim,
+            elementwise_affine=False,
+            eps=1e-6,
+            quant_config=quant_config,
+            prefix="norm_out",
+        )
         self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
 
     def forward(

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -113,7 +113,7 @@ class AsyncOmniDiffusion:
                 od_config.update_multimodal_support()
 
                 tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model)
-                od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+                od_config.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
             else:
                 raise FileNotFoundError("model_index.json not found")
         except (AttributeError, OSError, ValueError, FileNotFoundError):
@@ -121,7 +121,7 @@ class AsyncOmniDiffusion:
             if cfg is None:
                 raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
 
-            od_config.tf_model_config = TransformerConfig.from_dict(cfg)
+            od_config.set_tf_model_config(TransformerConfig.from_dict(cfg))
             model_type = cfg.get("model_type")
             architectures = cfg.get("architectures") or []
             # Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually

vllm_omni/quantization/factory.py

@@ -41,9 +41,25 @@ def _build_int8(**kw: Any) -> QuantizationConfig:
     return DiffusionInt8Config(**kw)
 
 
+def _build_inc(**kw: Any) -> QuantizationConfig:
+    """Lazy import for INC/AutoRound config with checkpoint kwarg normalization."""
+    from vllm.model_executor.layers.quantization.inc import INCConfig
+
+    # Map checkpoint key 'bits' to INCConfig's 'weight_bits'
+    if "bits" in kw and "weight_bits" not in kw:
+        kw["weight_bits"] = kw.pop("bits")
+
+    # Filter to only valid INCConfig params
+    valid = set(inspect.signature(INCConfig.__init__).parameters) - {"self"}
+    filtered = {k: v for k, v in kw.items() if k in valid}
+    return INCConfig(**filtered)
+
+
 _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
     "gguf": _build_gguf,
     "int8": _build_int8,
+    "inc": _build_inc,
+    "auto-round": _build_inc,
 }
 
 SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))