[Rebase] Align vLLM-Omni with vLLM v0.31.0 (#8459)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/diffusion/quantization/test_fp8_config.py

@@ -15,6 +15,74 @@ def test_build_quant_config_fp8():
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "dynamic"
+    assert not config.is_checkpoint_fp8_serialized
+
+
+def test_build_quant_config_fp8_uses_native_online_quantization():
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("fp8", ignored_layers=["proj_out"])
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args.linear.weight == kFp8StaticTensorSym
+    assert config.args.moe.weight == kFp8StaticTensorSym
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_build_quant_config_upstream_online_fp8_shorthand():
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("fp8_per_tensor", ignore=["proj_out"])
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args.linear.weight == kFp8StaticTensorSym
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_online_fp8_rejects_static_activation():
+    from vllm_omni.quantization import build_quant_config
+
+    with pytest.raises(ValueError, match="activation_scheme='dynamic'"):
+        build_quant_config("fp8", activation_scheme="static")
+
+
+def test_build_online_quant_config_preserves_explicit_native_args():
+    from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
+    from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+
+    from vllm_omni.quantization import build_quant_config
+
+    args = QuantizationConfigArgs(linear=QuantSpec(weight="fp8_per_tensor_static"), ignore=["proj_out"])
+    config = build_quant_config("online", args=args)
+    assert isinstance(config, OnlineQuantizationConfig)
+    assert config.args is args
+    assert config.ignored_layers == ["proj_out"]
+
+
+def test_build_online_quant_config_requires_args():
+    from vllm_omni.quantization import build_quant_config
+
+    with pytest.raises(ValueError, match="requires quantization config arguments"):
+        build_quant_config("online")
+
+
+def test_serialized_checkpoint_replaces_online_fp8_config():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.factory import resolve_quant_config_from_disk
+
+    config = resolve_quant_config_from_disk(
+        build_quant_config("fp8"),
+        {"quant_method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
+    )
+    assert isinstance(config, Fp8Config)
+    assert config.is_checkpoint_fp8_serialized
+    assert config.activation_scheme == "static"
 
 
 def test_build_quant_config_none():
@@ -39,7 +107,7 @@ def test_build_quant_config_invalid():
 def test_build_quant_config_dict():
     from vllm_omni.quantization import build_quant_config
 
-    config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
+    config = build_quant_config({"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"})
     assert config is not None
     assert config.get_name() == "fp8"
     assert config.activation_scheme == "static"
@@ -48,7 +116,7 @@ def test_build_quant_config_dict():
 def test_build_quant_config_dict_not_mutated():
     from vllm_omni.quantization import build_quant_config
 
-    original = {"method": "fp8", "activation_scheme": "static"}
+    original = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
     copy = original.copy()
     build_quant_config(original)
     assert original == copy
@@ -99,7 +167,7 @@ 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"}
+    inner = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
     original = inner.copy()
     build_quant_config({"transformer": inner, "vae": None})
     assert inner == original
@@ -115,20 +183,18 @@ def test_flat_dict_not_misdetected_as_per_component():
 
 
 def test_build_quant_config_passthrough():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     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
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     config = ComponentQuantizationConfig(
         component_configs={"transformer": fp8, "vae": None},
     )
@@ -140,11 +206,10 @@ def test_component_config_routing():
 
 
 def test_component_config_with_default():
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.quantization import ComponentQuantizationConfig
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    fp8 = DiffusionFp8Config()
     config = ComponentQuantizationConfig(
         component_configs={"vae": None},
         default_config=fp8,
@@ -167,7 +232,7 @@ def test_integration_dict():
 
     config = OmniDiffusionConfig(
         model="test",
-        quantization_config={"method": "fp8", "activation_scheme": "static"},
+        quantization_config={"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
     )
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "fp8"

tests/diffusion/quantization/test_online_fp8_gpu.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise online FP8 loading and inference through the native CUDA path."""
+
+import pytest
+import torch
+from torch.nn import functional as F
+from vllm.config import set_current_vllm_config
+from vllm.distributed import (
+    destroy_distributed_environment,
+    destroy_model_parallel,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
+from vllm.utils.network_utils import get_open_port
+from vllm.utils.torch_utils import set_default_torch_dtype
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
+from vllm_omni.quantization import build_quant_config
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize("method", ["fp8", "fp8_per_tensor"])
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_online_fp8_loads_bf16_weights_and_runs_native_kernel(method):
+    ignored_key = "ignored_layers" if method == "fp8" else "ignore"
+    config = build_quant_config(method, **{ignored_key: ["skip"]})
+    od_config = OmniDiffusionConfig(model="test", dtype=torch.bfloat16, quantization_config=config)
+    vllm_config = create_diffusion_vllm_config(torch.device("cuda:0"), od_config)
+
+    try:
+        init_distributed_environment(
+            world_size=1,
+            rank=0,
+            local_rank=0,
+            distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
+            backend="nccl",
+        )
+        with set_current_vllm_config(vllm_config), set_default_torch_dtype(torch.bfloat16), torch.device("cuda:0"):
+            initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
+            layer = ReplicatedLinear(
+                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="linear", disable_tp=True
+            )
+            assert isinstance(layer.quant_method, Fp8PerTensorOnlineLinearMethod)
+            assert layer.weight.is_meta
+
+            generator = torch.Generator(device="cuda").manual_seed(42)
+            weights = torch.randn(128, 128, generator=generator, dtype=torch.bfloat16) / 128**0.5
+            inputs = torch.randn(32, 128, generator=generator, dtype=torch.bfloat16)
+            expected = F.linear(inputs, weights).float()
+
+            # This is the loader installed by the native online quantizer.
+            # It materializes and quantizes the BF16 checkpoint in one pass.
+            layer.weight.weight_loader(layer.weight, weights)
+            assert layer.weight.dtype == torch.float8_e4m3fn
+            actual = layer(inputs)[0].float()
+            relative_mse = (actual - expected).square().mean() / expected.square().mean()
+            assert relative_mse.item() < 0.002
+
+            ignored = ReplicatedLinear(
+                128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="skip", disable_tp=True
+            )
+            assert isinstance(ignored.quant_method, UnquantizedLinearMethod)
+            assert ignored.weight.dtype == torch.bfloat16
+    finally:
+        destroy_model_parallel()
+        destroy_distributed_environment()