[Model] FP8 Support For Boogu-Image (#6925)

Signed-off-by: Dmaner <2663515256@qq.com>

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -330,6 +330,74 @@ def test_transformer_instantiates():
     assert model.x_embedder.out_features == HIDDEN_SIZE
 
 
+@pytest.fixture
+def recording_quant_config():
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    class RecordingQuantConfig(QuantizationConfig):
+        def __init__(self):
+            super().__init__()
+            self.prefixes: list[str] = []
+
+        @classmethod
+        def get_name(cls):
+            return "recording"
+
+        @classmethod
+        def get_supported_act_dtypes(cls):
+            return [torch.float32]
+
+        @classmethod
+        def get_min_capability(cls):
+            return 0
+
+        @classmethod
+        def get_config_filenames(cls):
+            return []
+
+        @classmethod
+        def from_config(cls, config):
+            return cls()
+
+        def get_quant_method(self, layer, prefix):
+            self.prefixes.append(prefix)
+            return UnquantizedLinearMethod()
+
+    return RecordingQuantConfig()
+
+
+def test_transformer_propagates_quant_config_and_prefix(recording_quant_config):
+    from vllm.model_executor.layers.linear import LinearBase
+
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageTransformer2DModel,
+    )
+
+    model = BooguImageTransformer2DModel(
+        od_config=_tiny_od_config(),
+        quant_config=recording_quant_config,
+        prefix="transformer",
+    )
+
+    configured_prefixes = set()
+    for name, module in model.named_modules():
+        if isinstance(module, LinearBase):
+            qualified_name = f"transformer.{name}"
+            assert module.prefix == qualified_name
+            assert module.quant_config is recording_quant_config
+            configured_prefixes.add(qualified_name)
+
+    assert set(recording_quant_config.prefixes) == configured_prefixes
+    assert {name for name, module in model.named_modules() if isinstance(module, torch.nn.Linear)} == {
+        "x_embedder",
+        "ref_image_patch_embedder",
+        "time_caption_embed.timestep_embedder.linear_1",
+        "time_caption_embed.timestep_embedder.linear_2",
+        "time_caption_embed.caption_embedder.1",
+    }
+
+
 def test_transformer_preprocesses_multiple_instruction_feature_layers():
     from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
         BooguImageTransformer2DModel,

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -83,6 +83,7 @@ def mock_dependencies(mocker, monkeypatch):
         "processor": mock_processor,
         "vae": mock_vae,
         "scheduler": mock_scheduler,
+        "transformer_cls": mock_transformer_cls,
         "transformer": mock_transformer_instance,
     }
 
@@ -187,6 +188,28 @@ def test_constructor_weights_sources(boogu_pipeline):
     assert source.fall_back_to_pt is True
 
 
+def test_constructor_routes_only_transformer_config(mock_dependencies, mocker):
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+    transformer_config = mocker.MagicMock(spec=QuantizationConfig)
+    encoder_config = mocker.MagicMock(spec=QuantizationConfig)
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config=ComponentQuantizationConfig(
+            {"transformer": transformer_config, "mllm": encoder_config, "vae": None}
+        ),
+    )
+    BooguImagePipeline(od_config=od_config)
+    kwargs = mock_dependencies["transformer_cls"].call_args.kwargs
+    assert kwargs["quant_config"] is transformer_config
+    assert kwargs["prefix"] == "transformer"
+
+
 @pytest.mark.parametrize(
     ("parallel_config", "cache_backend", "message"),
     [