[Core] Unified quantization framework (#1764)

tests/diffusion/quantization/test_fp8_config.py

@@ -1,137 +1,366 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for FP8 quantization config."""
+"""Tests for the unified quantization framework."""
 
 import pytest
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
-def test_fp8_config_creation():
-    """Test that FP8 config can be created."""
-    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+def test_build_quant_config_fp8():
+    from vllm_omni.quantization import build_quant_config
 
-    config = get_diffusion_quant_config("fp8")
+    config = build_quant_config("fp8")
     assert config is not None
     assert config.get_name() == "fp8"
+    assert config.activation_scheme == "dynamic"
 
 
-def test_vllm_config_extraction():
-    """Test that vLLM config can be extracted from diffusion config."""
-    from vllm_omni.diffusion.quantization import (
-        get_diffusion_quant_config,
-        get_vllm_quant_config_for_layers,
-    )
+def test_build_quant_config_none():
+    from vllm_omni.quantization import build_quant_config
 
-    diff_config = get_diffusion_quant_config("fp8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
-    assert vllm_config is not None
-    assert vllm_config.activation_scheme == "dynamic"
+    assert build_quant_config(None) is None
 
 
-def test_none_quantization():
-    """Test that None quantization returns None config."""
-    from vllm_omni.diffusion.quantization import (
-        get_diffusion_quant_config,
-        get_vllm_quant_config_for_layers,
-    )
+def test_build_quant_config_none_string():
+    from vllm_omni.quantization import build_quant_config
 
-    config = get_diffusion_quant_config(None)
-    assert config is None
-    vllm_config = get_vllm_quant_config_for_layers(config)
-    assert vllm_config is None
+    assert build_quant_config("none") is None
 
 
-def test_invalid_quantization():
-    """Test that invalid quantization method raises error."""
-    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+def test_build_quant_config_invalid():
+    from vllm_omni.quantization import build_quant_config
 
     with pytest.raises(ValueError, match="Unknown quantization method"):
-        get_diffusion_quant_config("invalid_method")
+        build_quant_config("invalid_method")
 
 
-def test_fp8_config_with_custom_params():
-    """Test FP8 config with custom parameters."""
-    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+def test_build_quant_config_dict():
+    from vllm_omni.quantization import build_quant_config
 
-    config = get_diffusion_quant_config(
-        "fp8",
-        activation_scheme="static",
-        ignored_layers=["proj_out"],
-    )
+    config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
     assert config is not None
+    assert config.get_name() == "fp8"
     assert config.activation_scheme == "static"
-    assert "proj_out" in config.ignored_layers
 
 
-def test_supported_methods():
-    """Test that supported methods list is correct."""
-    from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
+def test_build_quant_config_dict_not_mutated():
+    from vllm_omni.quantization import build_quant_config
 
-    assert "fp8" in SUPPORTED_QUANTIZATION_METHODS
+    original = {"method": "fp8", "activation_scheme": "static"}
+    copy = original.copy()
+    build_quant_config(original)
+    assert original == copy
 
 
-def test_quantization_integration():
-    """Test end-to-end quantization flow through OmniDiffusionConfig."""
+def test_build_quant_config_per_component():
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config(
+        {
+            "transformer": {"method": "fp8"},
+            "vae": None,
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+    assert config.component_configs["transformer"].get_name() == "fp8"
+    assert config.component_configs["vae"] is None
+
+
+def test_build_quant_config_per_component_string():
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config({"transformer": "fp8", "vae": None})
+    assert isinstance(config, ComponentQuantizationConfig)
+    assert config.component_configs["transformer"].get_name() == "fp8"
+
+
+def test_build_quant_config_per_component_inner_dict_not_mutated():
+    """Inner component dicts should not be mutated by build_quant_config."""
+    from vllm_omni.quantization import build_quant_config
+
+    inner = {"method": "fp8", "activation_scheme": "static"}
+    original = inner.copy()
+    build_quant_config({"transformer": inner, "vae": None})
+    assert inner == original
+
+
+def test_flat_dict_not_misdetected_as_per_component():
+    """A flat config like {"activation_scheme": "static"} must NOT be treated
+    as a per-component dict — it should raise ValueError for missing 'method'."""
+    from vllm_omni.quantization import build_quant_config
+
+    with pytest.raises(ValueError, match="must have a 'method' key"):
+        build_quant_config({"activation_scheme": "static"})
+
+
+def test_build_quant_config_passthrough():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import build_quant_config
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    assert build_quant_config(fp8) is fp8
+
+
+def test_component_config_routing():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    config = ComponentQuantizationConfig(
+        component_configs={"transformer": fp8, "vae": None},
+    )
+
+    assert config.get_name() == "component"
+    assert config.resolve("transformer.blocks.0.attn") is fp8
+    assert config.resolve("vae.encoder.conv_in") is None
+    assert config.resolve("unknown.layer") is None
+
+
+def test_component_config_with_default():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    config = ComponentQuantizationConfig(
+        component_configs={"vae": None},
+        default_config=fp8,
+    )
+
+    assert config.resolve("transformer.blocks.0") is fp8
+    assert config.resolve("vae.encoder") is None
+
+
+def test_gguf_config():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
+
+    config = build_quant_config(
+        {
+            "method": "gguf",
+            "gguf_model": "path/to/model.gguf",
+        }
+    )
+    assert isinstance(config, DiffusionGGUFConfig)
+    assert config.gguf_model == "path/to/model.gguf"
+    assert config.get_name() == "gguf"
+
+
+def test_integration_string():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    # Test with quantization string only
-    config = OmniDiffusionConfig(model="test", quantization="fp8")
+    config = OmniDiffusionConfig(model="test", quantization_config="fp8")
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "fp8"
 
-    # Test with quantization_config dict
-    config2 = OmniDiffusionConfig(
+
+def test_integration_dict():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig(
         model="test",
         quantization_config={"method": "fp8", "activation_scheme": "static"},
     )
-    assert config2.quantization_config is not None
-    assert config2.quantization_config.get_name() == "fp8"
-    assert config2.quantization_config.activation_scheme == "static"
-
-    # Test that vLLM config can be extracted
-    vllm_config = config.quantization_config.get_vllm_quant_config()
-    assert vllm_config is not None
-
-
-def test_quantization_dict_not_mutated():
-    """Test that passing a dict to quantization_config doesn't mutate it."""
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-    original_dict = {"method": "fp8", "activation_scheme": "static"}
-    dict_copy = original_dict.copy()
-
-    OmniDiffusionConfig(model="test", quantization_config=original_dict)
-
-    # Original dict should be unchanged
-    assert original_dict == dict_copy
-
-
-def test_quantization_conflicting_methods_warning(caplog):
-    """Test warning when quantization and quantization_config['method'] conflict."""
-    import logging
-
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-    with caplog.at_level(logging.WARNING):
-        config = OmniDiffusionConfig(
-            model="test",
-            quantization="fp8",  # This should be overridden
-            quantization_config={"method": "fp8", "activation_scheme": "static"},
-        )
-    # No warning when methods match
     assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "fp8"
+    assert config.quantization_config.activation_scheme == "static"
 
 
-def test_fp8_delegates_to_vllm_config():
-    """Test that DiffusionFp8Config delegates to vLLM's Fp8Config."""
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+def test_integration_no_quant():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    from vllm_omni.diffusion.quantization import DiffusionFp8Config
+    config = OmniDiffusionConfig(model="test")
+    assert config.quantization_config is None
 
-    # Test that quant_config_cls is set correctly
-    assert DiffusionFp8Config.quant_config_cls is Fp8Config
 
-    # Test that get_name() delegates to vLLM
-    assert DiffusionFp8Config.get_name() == Fp8Config.get_name()
+def test_integration_per_component():
+    """OmniDiffusionConfig with per-component quantization dict."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.quantization import ComponentQuantizationConfig
 
-    # Test that get_min_capability() delegates to vLLM
-    assert DiffusionFp8Config.get_min_capability() == Fp8Config.get_min_capability()
+    config = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "transformer": {"method": "fp8"},
+            "vae": None,
+        },
+    )
+    assert isinstance(config.quantization_config, ComponentQuantizationConfig)
+    assert config.quantization_config.component_configs["transformer"].get_name() == "fp8"
+    assert config.quantization_config.component_configs["vae"] is None
+
+
+def test_supported_methods_includes_vllm():
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt", "modelopt_fp4"]:
+        assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
+
+
+def test_supported_methods_count():
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert len(SUPPORTED_QUANTIZATION_METHODS) >= 20
+
+
+def test_per_component_routing_with_resolve():
+    """Verify resolve() routes correctly by prefix."""
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config(
+        {
+            "transformer": {"method": "fp8"},
+            "vae": None,
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+
+    assert config.resolve("transformer.blocks.0.attn.to_q") is not None
+    assert config.resolve("transformer.blocks.0.attn.to_q").get_name() == "fp8"
+    assert config.resolve("vae.encoder.conv_in") is None
+    assert config.resolve("unknown.layer.0.weight") is None
+
+
+def test_per_component_routing_with_default():
+    """Verify default config applies to unmatched prefixes."""
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config(
+        {
+            "vae": None,
+            "default": "fp8",
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+
+    assert config.resolve("vae.decoder.conv") is None
+    resolved = config.resolve("transformer.blocks.0.attn")
+    assert resolved is not None
+    assert resolved.get_name() == "fp8"
+
+
+@pytest.mark.parametrize(
+    "quant_algo",
+    ["FP8", "NVFP4"],
+    ids=["modelopt_fp8", "modelopt_nvfp4"],
+)
+def test_omni_convertor_thinker_finds_text_config_quant(quant_algo):
+    """Thinker stage should discover quantization_config from
+    thinker_config.text_config for both FP8 and NVFP4 modelopt checkpoints."""
+    from types import SimpleNamespace
+
+    from vllm_omni.config.model import OmniModelArchConfigConvertor
+
+    text_config = SimpleNamespace(
+        quantization_config={
+            "quant_method": "modelopt",
+            "quant_algo": quant_algo,
+            "ignore": ["lm_head", "model.layers.0.mlp.gate"],
+        },
+        model_type="qwen3_moe",
+    )
+    thinker_config = SimpleNamespace(text_config=text_config)
+    hf_config = SimpleNamespace(
+        thinker_config=thinker_config,
+        talker_config=SimpleNamespace(text_config=SimpleNamespace()),
+        model_type="qwen3_omni_moe",
+    )
+
+    convertor = OmniModelArchConfigConvertor(hf_config, text_config, stage_config_name="thinker_config")
+    quant_cfg = convertor.get_quantization_config()
+
+    assert quant_cfg is not None
+    assert quant_cfg["quant_method"] == "modelopt"
+    assert "lm_head" in quant_cfg["ignore"]
+
+
+def test_omni_convertor_talker_returns_none():
+    """Talker stage should get no quantization config when its text_config
+    has no quantization_config (talker weights are BF16)."""
+    from types import SimpleNamespace
+
+    from vllm_omni.config.model import OmniModelArchConfigConvertor
+
+    talker_text_config = SimpleNamespace(model_type="qwen3_omni_moe_talker")
+    talker_config = SimpleNamespace(text_config=talker_text_config)
+    hf_config = SimpleNamespace(
+        talker_config=talker_config,
+        model_type="qwen3_omni_moe",
+    )
+
+    convertor = OmniModelArchConfigConvertor(hf_config, talker_text_config, stage_config_name="talker_config")
+    quant_cfg = convertor.get_quantization_config()
+
+    assert quant_cfg is None
+
+
+def test_omni_convertor_no_stage_name_falls_back():
+    """Without stage_config_name, should fall back to base behavior."""
+    from types import SimpleNamespace
+
+    from vllm_omni.config.model import OmniModelArchConfigConvertor
+
+    hf_config = SimpleNamespace(model_type="qwen3_omni_moe")
+    text_config = SimpleNamespace()
+
+    convertor = OmniModelArchConfigConvertor(hf_config, text_config)
+    quant_cfg = convertor.get_quantization_config()
+
+    assert quant_cfg is None
+
+
+def test_multi_component_model_routing():
+    """Integration test: walk a multi-component model and verify per-component
+    quantization routes resolve() correctly for every linear layer."""
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    # Build a mock multi-stage model mimicking Bagel/Qwen3-Omni layout
+    class MockTransformerBlock(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.attn_q = nn.Linear(64, 64)
+            self.attn_k = nn.Linear(64, 64)
+            self.mlp = nn.Linear(64, 256)
+
+    class MockVAEBlock(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = nn.Linear(64, 64)
+
+    class MockMultiStageModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.ModuleDict(
+                {
+                    "block_0": MockTransformerBlock(),
+                    "block_1": MockTransformerBlock(),
+                }
+            )
+            self.vae = nn.ModuleDict(
+                {
+                    "encoder": MockVAEBlock(),
+                    "decoder": MockVAEBlock(),
+                }
+            )
+
+    model = MockMultiStageModel()
+    config = build_quant_config(
+        {
+            "transformer": {"method": "fp8", "activation_scheme": "dynamic"},
+            "vae": None,
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear):
+            resolved = config.resolve(name)
+            if name.startswith("transformer"):
+                assert resolved is not None, f"{name} should be quantized"
+                assert resolved.get_name() == "fp8"
+            elif name.startswith("vae"):
+                assert resolved is None, f"{name} should NOT be quantized"

tests/diffusion/quantization/test_gguf_config.py

@@ -9,16 +9,12 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm.model_executor.layers.linear import LinearBase
-from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES
+from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES, UnquantizedLinearMethod
 
-from vllm_omni.diffusion.quantization import (
-    DiffusionGgufConfig,
-    get_diffusion_quant_config,
-)
-from vllm_omni.diffusion.quantization.gguf import (
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.gguf_config import (
+    DiffusionGGUFConfig,
     DiffusionGGUFLinearMethod,
-    UnquantizedLinearMethod,
-    _GGUFConfig,
     dequant_gemm_gguf,
 )
 
@@ -26,7 +22,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def test_gguf_config_creation_and_delegation():
-    config = DiffusionGgufConfig(
+    config = DiffusionGGUFConfig(
         gguf_model="weights.gguf",
         unquantized_modules=["proj_out"],
     )
@@ -34,37 +30,36 @@ def test_gguf_config_creation_and_delegation():
     assert config.gguf_model == "weights.gguf"
     assert config.unquantized_modules == ["proj_out"]
     assert config.get_name() == "gguf"
-    assert isinstance(config.get_vllm_quant_config(), _GGUFConfig)
 
 
-def test_get_diffusion_quant_config_builds_gguf_config():
-    config = get_diffusion_quant_config(
+def test_build_quant_config_builds_gguf_config():
+    config = build_quant_config(
         "gguf",
         gguf_model="repo/model-Q4_0.gguf",
         unquantized_modules=["norm_out"],
     )
 
-    assert isinstance(config, DiffusionGgufConfig)
+    assert isinstance(config, DiffusionGGUFConfig)
     assert config.gguf_model == "repo/model-Q4_0.gguf"
     assert config.unquantized_modules == ["norm_out"]
 
 
 def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
     linear = object.__new__(LinearBase)
-    method = _GGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
+    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
 
     assert isinstance(method, DiffusionGGUFLinearMethod)
 
 
 def test_gguf_vllm_config_respects_unquantized_modules():
     linear = object.__new__(LinearBase)
-    method = _GGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
+    method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
 
     assert isinstance(method, UnquantizedLinearMethod)
 
 
 def test_gguf_vllm_config_returns_none_for_non_linear_layers():
-    method = _GGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
+    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
 
     assert method is None
 
@@ -81,7 +76,7 @@ def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
 
 def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
     qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[]))
+    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
     layer = SimpleNamespace(
         qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
         qweight_type=SimpleNamespace(weight_type=qweight_type),
@@ -97,7 +92,7 @@ def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
 
 def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
     qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[]))
+    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
     qweight = torch.nn.Parameter(
         torch.tensor(
             [

tests/diffusion/quantization/test_int8_config.py

@@ -10,11 +10,11 @@ from pytest_mock import MockerFixture
 from torch.nn import Module, Parameter
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
 
-from vllm_omni.diffusion.quantization import (
-    get_diffusion_quant_config,
-    get_vllm_quant_config_for_layers,
-)
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.")
 
@@ -23,36 +23,26 @@ cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason=
 
 def test_int8_config_creation():
     """Test that Int8 config can be created."""
-    config = get_diffusion_quant_config("int8")
+    config = build_quant_config("int8")
     assert config is not None
     assert config.get_name() == "int8"
 
 
-def test_vllm_config_extraction():
-    """Test that vLLM config can be extracted from diffusion config."""
-    diff_config = get_diffusion_quant_config("int8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
-    assert vllm_config is not None
-    assert vllm_config.activation_scheme == "dynamic"
-
-
 def test_none_quantization():
     """Test that None quantization returns None config."""
-    config = get_diffusion_quant_config(None)
+    config = build_quant_config(None)
     assert config is None
-    vllm_config = get_vllm_quant_config_for_layers(config)
-    assert vllm_config is None
 
 
 def test_invalid_quantization():
     """Test that invalid quantization method raises error."""
     with pytest.raises(ValueError, match="Unknown quantization method"):
-        get_diffusion_quant_config("invalid_method")
+        build_quant_config("invalid_method")
 
 
 def test_int8_config_with_custom_params():
     """Test Int8 config with custom parameters."""
-    config = get_diffusion_quant_config(
+    config = build_quant_config(
         "int8",
         activation_scheme="dynamic",
         ignored_layers=["proj_out"],
@@ -64,8 +54,6 @@ def test_int8_config_with_custom_params():
 
 def test_supported_methods():
     """Test that supported methods list is correct."""
-    from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
-
     assert "int8" in SUPPORTED_QUANTIZATION_METHODS
 
 
@@ -73,8 +61,8 @@ def test_quantization_integration():
     """Test end-to-end quantization flow through OmniDiffusionConfig."""
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    # Test with quantization string only
-    config = OmniDiffusionConfig(model="test", quantization="int8")
+    # Test with quantization_config string
+    config = OmniDiffusionConfig(model="test", quantization_config="int8")
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "int8"
 
@@ -87,10 +75,6 @@ def test_quantization_integration():
     assert config2.quantization_config.get_name() == "int8"
     assert config2.quantization_config.activation_scheme == "dynamic"
 
-    # Test that vLLM config can be extracted
-    vllm_config = config.quantization_config.get_vllm_quant_config()
-    assert vllm_config is not None
-
 
 def test_quantization_dict_not_mutated():
     """Test that passing a dict to quantization_config doesn't mutate it."""
@@ -105,28 +89,24 @@ def test_quantization_dict_not_mutated():
     assert original_dict == dict_copy
 
 
-def test_quantization_conflicting_methods_warning(caplog):
-    """Test warning when quantization and quantization_config['method'] conflict."""
-    import logging
-
+def test_quantization_config_string_and_dict_equivalent():
+    """Test that string and dict forms produce equivalent configs."""
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    with caplog.at_level(logging.WARNING):
-        config = OmniDiffusionConfig(
-            model="test",
-            quantization="int8",  # This should be overridden
-            quantization_config={"method": "int8", "activation_scheme": "dynamic"},
-        )
-    # No warning when methods match
-    assert config.quantization_config is not None
+    config_str = OmniDiffusionConfig(model="test", quantization_config="int8")
+    config_dict = OmniDiffusionConfig(
+        model="test",
+        quantization_config={"method": "int8", "activation_scheme": "dynamic"},
+    )
+    assert config_str.quantization_config.get_name() == config_dict.quantization_config.get_name()
+    assert config_str.quantization_config.activation_scheme == config_dict.quantization_config.activation_scheme
 
 
 def test_get_quant_method(mocker: MockerFixture):
     """Test for get_quant_method method for GPU"""
-    from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+    from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
 
-    diff_config = get_diffusion_quant_config("int8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+    config = build_quant_config("int8")
 
     def _fake_init(self, quant_config):
         pass
@@ -141,21 +121,20 @@ def test_get_quant_method(mocker: MockerFixture):
         patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True),
         patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False),
     ):
-        method = vllm_config.get_quant_method(layer, prefix)
+        method = config.get_quant_method(layer, prefix)
         assert isinstance(method, Int8OnlineLinearMethod)
 
     # Test skipping quantization for a layer
-    vllm_config.ignored_layers = [prefix]
-    method = vllm_config.get_quant_method(layer, prefix)
+    config.ignored_layers = [prefix]
+    method = config.get_quant_method(layer, prefix)
     assert isinstance(method, UnquantizedLinearMethod)
 
 
 def test_get_npu_quant_method():
     """Test for get_quant_method method for NPU"""
-    from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
 
-    diff_config = get_diffusion_quant_config("int8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+    config = build_quant_config("int8")
 
     layer = MagicMock(spec=LinearBase)
     prefix = "test_layer"
@@ -165,12 +144,12 @@ def test_get_npu_quant_method():
         patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False),
         patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True),
     ):
-        method = vllm_config.get_quant_method(layer, prefix)
+        method = config.get_quant_method(layer, prefix)
         assert isinstance(method, NPUInt8OnlineLinearMethod)
 
     # Test skipping quantization for a layer
-    vllm_config.ignored_layers = [prefix]
-    method = vllm_config.get_quant_method(layer, prefix)
+    config.ignored_layers = [prefix]
+    method = config.get_quant_method(layer, prefix)
     assert isinstance(method, UnquantizedLinearMethod)
 
 
@@ -189,12 +168,12 @@ class TestInt8LinearMethod:
     @pytest.fixture
     def patch_deps(self, mocker, mock_kernel):
         # mock init_int8_linear_kernel
-        mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
+        mocker.patch("vllm_omni.quantization.int8_config.init_int8_linear_kernel", return_value=mock_kernel)
         return mock_kernel
 
     def test_init(self, patch_deps, mock_quant_config):
         # test for Int8LinearMethod init
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod, init_int8_linear_kernel
+        from vllm_omni.quantization.int8_config import Int8LinearMethod, init_int8_linear_kernel
 
         method = Int8LinearMethod(mock_quant_config)
 
@@ -205,7 +184,7 @@ class TestInt8LinearMethod:
         assert method.int8_linear == patch_deps
 
     def test_process_weights_after_loading(self, patch_deps, mock_quant_config):
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+        from vllm_omni.quantization.int8_config import Int8LinearMethod
 
         method = Int8LinearMethod(mock_quant_config)
         layer = Module()
@@ -214,7 +193,7 @@ class TestInt8LinearMethod:
         patch_deps.process_weights_after_loading.assert_called_once_with(layer)
 
     def test_apply(self, patch_deps, mock_quant_config):
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+        from vllm_omni.quantization.int8_config import Int8LinearMethod
 
         method = Int8LinearMethod(mock_quant_config)
         layer = Module()
@@ -237,19 +216,19 @@ class TestInt8OnlineLinearMethod:
         # mock kernel
         mock_kernel = mocker.Mock()
         mock_kernel.layer_param_names = ("weight", "weight_scale", "input_scale", "input_zero_point", "azp_adj")
-        mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
-        mocker.patch("vllm_omni.diffusion.quantization.int8.replace_parameter")
+        mocker.patch("vllm_omni.quantization.int8_config.init_int8_linear_kernel", return_value=mock_kernel)
+        mocker.patch("vllm_omni.quantization.int8_config.replace_parameter")
 
         # mock scaled_int8_quant return value
         mock_qweight = torch.ones((128, 64), dtype=torch.int8)
         mock_scale = torch.randn(128)
         mock_quant = mocker.patch(
-            "vllm_omni.diffusion.quantization.int8.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
+            "vllm_omni.quantization.int8_config.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
         )
         return {"kernel": mock_kernel, "quant": mock_quant, "mock_qweight": mock_qweight, "mock_scale": mock_scale}
 
     def test_process_weights_after_loading(self, mock_deps, mock_quant_config):
-        from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
 
         method = Int8OnlineLinearMethod(mock_quant_config)
         layer = Module()
@@ -268,12 +247,12 @@ class TestNPUInt8LinearMethod:
     def mock_torch_npu(self, mocker):
         torch_npu = MagicMock()
 
-        mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu", return_value=torch_npu)
+        mocker.patch("vllm_omni.quantization.int8_config.torch_npu", return_value=torch_npu)
         mocker.patch(
-            "vllm_omni.diffusion.quantization.int8.torch_npu.npu_dynamic_quant",
+            "vllm_omni.quantization.int8_config.torch_npu.npu_dynamic_quant",
             return_value=(self.qweight_mock, self.scale_mock),
         )
-        mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu.npu_quant_matmul", return_value=self.out_mock)
+        mocker.patch("vllm_omni.quantization.int8_config.torch_npu.npu_quant_matmul", return_value=self.out_mock)
         return torch_npu
 
     @pytest.fixture
@@ -288,7 +267,7 @@ class TestNPUInt8LinearMethod:
         return layer
 
     def test_npu_int8_process_weights_after_loading(self, mock_layer, mock_quant_config, mock_torch_npu):
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
 
         method = NPUInt8LinearMethod(mock_quant_config)
         ori_weight_shape = mock_layer.weight.shape
@@ -299,7 +278,7 @@ class TestNPUInt8LinearMethod:
         assert mock_layer.weight.is_contiguous()
 
     def test_npu_int8_apply(self, mock_layer, mock_quant_config, mock_torch_npu):
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
 
         method = NPUInt8LinearMethod(mock_quant_config)
         x = torch.randn(1, 16, 64)
@@ -308,7 +287,7 @@ class TestNPUInt8LinearMethod:
         assert output.shape == (1, 16, 128)
 
     def test_npu_int8_online_process_weights(self, mock_layer, mock_quant_config, mock_torch_npu):
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
 
         method = NPUInt8OnlineLinearMethod(mock_quant_config)
         method.process_weights_after_loading(mock_layer)
@@ -320,9 +299,9 @@ class TestNPUInt8LinearMethod:
 @pytest.fixture
 def quant_config():
     """Shared quant config fixture for smoke tests."""
-    from vllm_omni.diffusion.quantization.int8 import Int8Config
+    from vllm_omni.quantization.int8_config import DiffusionInt8Config
 
-    return Int8Config(
+    return DiffusionInt8Config(
         is_checkpoint_int8_serialized=False,
         activation_scheme="dynamic",
     )
@@ -360,7 +339,7 @@ class TestNPUInt8Smoke:
 
     def test_real_npu_online_process_weights_after_loading(self, quant_config, real_layer):
         """Smoke test: full process_weights_after_loading with real torch_npu."""
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
 
         method = NPUInt8OnlineLinearMethod(quant_config)
 
@@ -375,7 +354,7 @@ class TestNPUInt8Smoke:
         """Smoke test: forward pass with real npu_quant_matmul."""
         import torch_npu
 
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
 
         method = NPUInt8LinearMethod(quant_config)
 
@@ -425,7 +404,7 @@ class TestCudaInt8Smoke:
 
     def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
         """Smoke test: full process_weights_after_loading with real CUDA ops."""
-        from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
 
         method = Int8OnlineLinearMethod(quant_config)
 
@@ -439,7 +418,7 @@ class TestCudaInt8Smoke:
         """Smoke test: forward pass with real CUDA int8 kernel."""
         from vllm import _custom_ops as ops
 
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+        from vllm_omni.quantization.int8_config import Int8LinearMethod
 
         method = Int8LinearMethod(quant_config)
 

tests/diffusion/quantization/test_quantization_quality.py

@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Quantization quality gate — validates that quantized models stay within
+acceptable perceptual distance (LPIPS) of the BF16 baseline.
+
+Developers adding a new quantization method should:
+1. Add their method + model to QUALITY_CONFIGS below
+2. Set a max_lpips threshold (use 0.15 for image, 0.20 for video as defaults)
+3. Run: pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
+4. Paste the output table into their PR description
+
+The test generates outputs with both BF16 and the quantized method using the
+same seed, computes LPIPS, and fails if it exceeds the threshold.
+
+Requirements:
+    pip install lpips
+
+Example — run only FP8 tests:
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "fp8"
+
+Example — run a specific model:
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
+"""
+
+from __future__ import annotations
+
+import gc
+from dataclasses import dataclass
+
+import numpy as np
+import pytest
+import torch
+
+from tests.utils import hardware_marks
+
+# ---------------------------------------------------------------------------
+# Configuration — add new quantization methods / models here
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class QualityTestConfig:
+    """Defines a single quantization quality test case."""
+
+    id: str  # pytest ID, e.g. "fp8_z_image"
+    model: str  # HF model name
+    quantization: str  # quantization method, e.g. "fp8"
+    task: str  # "t2i" or "t2v"
+    prompt: str  # generation prompt
+    max_lpips: float  # fail threshold — higher = more lenient
+    height: int = 1024
+    width: int = 1024
+    num_inference_steps: int = 20  # keep low for CI speed
+    num_frames: int = 5  # only for t2v
+    seed: int = 42
+    gpu: str = "H100"  # minimum GPU requirement
+
+
+# Add new quantization methods / models here.
+# Developers: copy a config, change quantization + max_lpips, run the test.
+QUALITY_CONFIGS = [
+    QualityTestConfig(
+        id="fp8_z_image",
+        model="Tongyi-MAI/Z-Image-Turbo",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table, morning light",
+        max_lpips=0.10,
+        num_inference_steps=20,
+    ),
+    QualityTestConfig(
+        id="fp8_flux",
+        model="black-forest-labs/FLUX.1-dev",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table, morning light",
+        max_lpips=0.20,
+        num_inference_steps=10,
+    ),
+    QualityTestConfig(
+        id="fp8_qwen_image",
+        model="Qwen/Qwen-Image",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table, morning light",
+        max_lpips=0.35,
+        seed=142,
+        num_inference_steps=20,
+    ),
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _generate_image(omni, config: QualityTestConfig):
+    """Generate a single image, return (PIL.Image, peak_mem_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(
+        device=current_omni_platform.device_type,
+    ).manual_seed(config.seed)
+    torch.cuda.reset_peak_memory_stats()
+
+    outputs = omni.generate(
+        {"prompt": config.prompt},
+        OmniDiffusionSamplingParams(
+            height=config.height,
+            width=config.width,
+            generator=generator,
+            num_inference_steps=config.num_inference_steps,
+        ),
+    )
+
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    first = outputs[0]
+    req_out = first.request_output[0] if hasattr(first, "request_output") else first
+    return req_out.images[0], peak_mem
+
+
+def _generate_video(omni, config: QualityTestConfig):
+    """Generate a video, return (np.ndarray [F,H,W,C], peak_mem_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(
+        device=current_omni_platform.device_type,
+    ).manual_seed(config.seed)
+    torch.cuda.reset_peak_memory_stats()
+
+    outputs = omni.generate(
+        {"prompt": config.prompt, "negative_prompt": ""},
+        OmniDiffusionSamplingParams(
+            height=config.height,
+            width=config.width,
+            generator=generator,
+            num_inference_steps=config.num_inference_steps,
+            num_frames=config.num_frames,
+        ),
+    )
+
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    first = outputs[0]
+    if hasattr(first, "request_output") and isinstance(first.request_output, list):
+        inner = first.request_output[0]
+        if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
+            frames = inner.images[0] if inner.images else None
+        else:
+            frames = inner
+    elif hasattr(first, "images") and first.images:
+        frames = first.images
+    else:
+        raise ValueError("Could not extract video frames from output.")
+
+    if isinstance(frames, torch.Tensor):
+        video = frames.detach().cpu()
+        if video.dim() == 5:
+            video = video[0]
+        if video.dim() == 4 and video.shape[0] in (3, 4):
+            video = video.permute(1, 2, 3, 0)
+        if video.is_floating_point():
+            video = video.clamp(-1, 1) * 0.5 + 0.5
+        return video.float().numpy(), peak_mem
+
+    return np.asarray(frames), peak_mem
+
+
+def _compute_lpips(baseline, quantized, task: str) -> float:
+    """Compute LPIPS between baseline and quantized outputs."""
+    from benchmarks.diffusion.quantization_quality import (
+        compute_lpips_images,
+        compute_lpips_video,
+    )
+
+    if task == "t2i":
+        return compute_lpips_images([baseline], [quantized])[0]
+    return compute_lpips_video(baseline, quantized)
+
+
+def _unload(omni):
+    del omni
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+_marks = hardware_marks(res={"cuda": "H100"})
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "config",
+    [pytest.param(c, id=c.id, marks=_marks) for c in QUALITY_CONFIGS],
+)
+def test_quantization_quality(config: QualityTestConfig):
+    """Validate that quantized output stays within LPIPS threshold of BF16."""
+    from vllm_omni.entrypoints.omni import Omni
+
+    generate_fn = _generate_video if config.task == "t2v" else _generate_image
+
+    # --- BF16 baseline ---
+    omni_bl = Omni(model=config.model)
+    baseline_out, bl_mem = generate_fn(omni_bl, config)
+    _unload(omni_bl)
+
+    # --- Quantized ---
+    omni_qt = Omni(model=config.model, quantization_config=config.quantization)
+    quant_out, qt_mem = generate_fn(omni_qt, config)
+    _unload(omni_qt)
+
+    # --- LPIPS ---
+    lpips_score = _compute_lpips(baseline_out, quant_out, config.task)
+
+    # --- Report ---
+    mem_reduction = (bl_mem - qt_mem) / bl_mem * 100 if bl_mem > 0 else 0
+    print(f"\n{'=' * 60}")
+    print(f"Quantization Quality: {config.id}")
+    print(f"{'=' * 60}")
+    print(f"  Model:         {config.model}")
+    print(f"  Method:        {config.quantization}")
+    print(f"  LPIPS:         {lpips_score:.4f}  (threshold: {config.max_lpips})")
+    print(f"  BF16 memory:   {bl_mem:.2f} GiB")
+    print(f"  Quant memory:  {qt_mem:.2f} GiB  ({mem_reduction:.0f}% reduction)")
+    print(f"  Result:        {'PASS' if lpips_score <= config.max_lpips else 'FAIL'}")
+    print(f"{'=' * 60}\n")
+
+    assert lpips_score <= config.max_lpips, (
+        f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} for {config.quantization} on {config.model}"
+    )