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

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

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -239,11 +239,11 @@ def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mo
 )
 def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
     from transformers import AutoModel
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
     from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym
 
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     od_config = OmniDiffusionConfig(
         model="dummy-boogu",
@@ -287,15 +287,14 @@ def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_co
         mock_dependencies["mllm_config_loader"].assert_not_called()
         mock_dependencies["mllm_builder"].assert_not_called()
         assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
-    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config)
+    assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config)
 
 
 def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
-    quant_config = Fp8Config(
+    quant_config = DiffusionFp8Config(
         ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
     )
     od_config = OmniDiffusionConfig(

tests/diffusion/models/flux2/test_flux2_quantization.py

@@ -16,17 +16,13 @@ class _RecordingBlock(nn.Module):
 
 
 def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch):
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
     from vllm_omni.diffusion.models.flux2 import flux2_transformer
+    from vllm_omni.quantization import build_quant_config
 
     monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
     monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)
 
-    quant_config = Fp8Config(
-        is_checkpoint_fp8_serialized=False,
-        activation_scheme="dynamic",
-    )
+    quant_config = build_quant_config("fp8")
     model = flux2_transformer.Flux2Transformer2DModel(
         in_channels=4,
         out_channels=4,

tests/diffusion/models/minimax_h3/test_minimax_h3_quantization.py

@@ -71,12 +71,12 @@ def _small_od_config(*, class_name=None):
 
 def test_fp8_scope_and_prefix_propagation(monkeypatch):
     from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
     from vllm.model_executor.layers.quantization.utils.quant_utils import (
         is_layer_skipped,
     )
 
     from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+    from vllm_omni.quantization.fp8_config import DiffusionFp8Config
 
     monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
     monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
@@ -92,7 +92,7 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
         "blocks.0.adaln_proj.linear",
         "final_layer.adaln_proj.linear",
     }
-    fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
+    fp8_config = DiffusionFp8Config(ignored_layers=sorted(ignored_layers))
     model = h3.MiniMaxH3DiTModel(
         _small_od_config(),
         quant_config=fp8_config,