[Model] Add online FP8 support for Boogu-Image MLLM (#7342)

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

recipes/Boogu/Boogu-Image.md

@@ -132,18 +132,19 @@ with `guidance_scale=1.0` remain valid on the same server: every rank evaluates
 only the positive branch, no negative embeddings are built, and no guidance is
 applied.
 
-### 1 x RTX 4090 48 GB (FP8)
+### 1 x RTX 4090 24 GB / 48 GB (FP8)
 
-Boogu-Image supports two ways to quantize or load the DiT transformer:
+Boogu-Image supports the following FP8 loading paths:
 
-| Path               | Base model                       | Edit model                       | Method                       | Scheme                  |
-| ------------------ | -------------------------------- | -------------------------------- | ---------------------------- | ----------------------- |
-| Serialized TorchAO | `Boogu/Boogu-Image-0.1-Base-fp8` | `Boogu/Boogu-Image-0.1-Edit-fp8` | `torchao_float8_weight_only` | FP8 weight-only (W8A16) |
-| Native online FP8  | `Boogu/Boogu-Image-0.1-Base`     | `Boogu/Boogu-Image-0.1-Edit`     | `fp8`                        | Dynamic W8A8            |
+| Path | Checkpoints | MLLM | DiT |
+| --- | --- | --- | --- |
+| Pre-quantized FP8 | `Boogu/Boogu-Image-0.1-{Base,Edit}-fp8` | HF checkpoint FP8 configuration | TorchAO FP8 weight-only (W8A16) |
+| Online FP8 | `Boogu/Boogu-Image-0.1-{Base,Edit}` | Optional 128x128 block FP8: Triton W8A8 (vLLM 0.28.0); Marlin W8A16 (vLLM 0.29.0) | Native dynamic FP8 (CUTLASS W8A8) |
 
 #### Environment
 
-The following configuration was validated:
+The following isolated environments were used for Base T2I and Edit I2I
+measurements at 512x512:
 
 - OS: Linux
 - Python: 3.12
@@ -152,7 +153,9 @@ The following configuration was validated:
 - vLLM-Omni version or commit: Use a commit that contains TorchAO FP8
   checkpoint loading for diffusion models
 - TorchAO version: 0.17.0
-- `kernels` / `kernels-data`: 0.16.1 / matching kernels-data
+- `kernels` / `kernels-data`: 0.16.1 / 0.16.1
+- Transformers: 5.17.0
+- Tokenizers: 0.23.1
 
 #### Command
 
@@ -167,19 +170,61 @@ vllm serve Boogu/Boogu-Image-0.1-Base-fp8 \
   '{"transformer":{"method":"torchao_float8_weight_only"}}'
 ```
 
-Native online FP8 (currently applied only to eligible vLLM-quantizable linear
-layers in the DiT transformer):
+Native online FP8:
 
 ```bash
+# DiT only
 vllm serve Boogu/Boogu-Image-0.1-Base \
   --omni \
   --port 8091 \
   --diffusion-quantization-config \
-  '{"transformer":{"method":"fp8"}}'
+  '{"mllm":null,"transformer":"fp8","vae":null}'
+
+# MLLM + DiT
+vllm serve Boogu/Boogu-Image-0.1-Base \
+  --omni \
+  --port 8091 \
+  --quantization fp8
+
+# MLLM + DiT (using --diffusion-quantization-config)
+vllm serve Boogu/Boogu-Image-0.1-Base \
+  --omni \
+  --port 8091 \
+  --diffusion-quantization-config \
+  '{"mllm":"fp8","transformer":"fp8","vae":null}'
+
 ```
 
-For image editing, use the matching Edit model ID from the table and keep the
-corresponding quantization configuration.
+For image editing, replace `Base` with `Edit` and keep the corresponding
+quantization configuration.
+
+#### Notes
+
+- **Memory usage:**
+    - The serialized TorchAO FP8 path uses approximately 21.1 GiB at 512x512
+      and 24.2-24.6 GiB at 1024x1024.
+    - The DiT-only online FP8 path uses approximately 26.9 GiB at 512x512
+      and 29.9-30.9 GiB at 1024x1024.
+    - The MLLM + DiT online FP8 path uses approximately 20.5-20.6 GiB at
+      512x512 and 23.6-24.5 GiB at 1024x1024.
+
+- **Quantization scope:** For original Base/Edit checkpoints, `--quantization fp8`
+  quantizes the MLLM's language-model linear layers through per-block FP8
+  and the DiT's linear layers through native online FP8. The vision encoder,
+  embeddings, normalization layers and VAE remain unquantized.
+
+- **Performance:** In the 48 GB Base/Edit benchmarks at 512x512, both DiT-only
+  and MLLM + DiT online FP8 reduced E2E latency by over 10 seconds compared with
+  official FP8 and vLLM-Omni pre-quantized FP8. The main runtime difference lies in the DiT quantization implementation, which switches from TorchAO weight-only FP8 (W8A16) to native CUTLASS
+  W8A8 kernels. This reduces additional dequantization overhead and enables
+  FP8 Tensor Core acceleration.
+
+- **Known Limitation:**
+    - **Output consistency:** With MLLM + DiT online FP8, same-seed outputs may visibly differ from BF16 for some Base prompts, but these differences do not mean lower image quality.
+
+    - **Import compatibility conflict between TorchAO and Diffusers::** Loading pre-quantized checkpoints requires `torchao>=0.17.0`, which has a compatibility conflict with `diffusers==0.40.0`. This produces a warning when Diffusers is imported, but currently does not affect Boogu-Image weight loading.
+
+    - **Conflict between the default kernels and Transformers versions:** The default Transformers version range (`>=5.10.1,<5.15`) conflicts with `kernels==0.16.1` when loading pre-quantized FP8 MLLM weights through Hugging Face `FP8Linear`. When using this loading path with `kernels==0.16.1`, please upgrade Transformers to `5.17.0`.
 
 #### Verification
 
@@ -190,19 +235,6 @@ ID passed to `vllm serve`. For example, use
 corresponding Edit model.
 
 Base text-to-image:
-## Fast text-to-image (Boogu-Image-0.1-Turbo)
-
-Turbo is the distilled text-to-image checkpoint. Its `model_index.json`
-declares `BooguImageTurboPipeline`, which runs the four-step DMD path instead
-of the scheduler-driven path used by Base.
-
-### Command
-
-```bash
-vllm serve Boogu/Boogu-Image-0.1-Turbo --omni --port 8091
-```
-
-### Verification
 
 ```bash
 curl -X POST http://localhost:8091/v1/images/generations \
@@ -231,22 +263,24 @@ curl -X POST http://localhost:8091/v1/images/edits \
   | jq -r '.data[0].b64_json' | base64 -d > edited_fp8.png
 ```
 
-#### Notes
+## Fast text-to-image (Boogu-Image-0.1-Turbo)
 
-- **Memory usage:**
-    - The serialized TorchAO FP8 path uses approximately 21.1 GiB at 512x512
-      and 24.2-24.6 GiB at 1024x1024.
-    - The native online FP8 path uses approximately 31.5 GiB for Base and
-      30.7 GiB for Edit at 1024x1024. It uses more memory because only eligible
-      DiT linear layers are quantized online, while the MLLM from the regular
-      Base/Edit checkpoint remains in BF16. Most of the observed difference
-      comes from the MLLM precision, although differences in the two DiT FP8
-      representations and kernels may also contribute.
-- **Quantization scope:** The transformer-scoped configuration applies to the
-  DiT. With the serialized `-fp8` checkpoints, the MLLM follows its own
-  checkpoint-declared FP8 configuration. With the regular checkpoints used for
-  online FP8, the MLLM remains in BF16. The VAE and scheduler are outside this
-  quantization routing.
+Turbo is the distilled text-to-image checkpoint. Its `model_index.json`
+declares `BooguImageTurboPipeline`, which runs the four-step DMD path instead
+of the scheduler-driven path used by Base.
+
+### Command
+
+```bash
+vllm serve Boogu/Boogu-Image-0.1-Turbo --omni --port 8091
+```
+
+### Verification
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
     "model": "Boogu/Boogu-Image-0.1-Turbo",
     "prompt": "A mountain lake at sunset, photorealistic, cinematic lighting",
     "size": "1024x1024",

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -13,11 +13,13 @@ Two groups:
    CFG handling, and reshape logic can be verified numerically on CPU.
 """
 
+import os
 from types import SimpleNamespace
 
 import pytest
 import torch
 from torch import nn
+from transformers import Qwen3VLConfig
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig, TransformerConfig
 
@@ -56,9 +58,17 @@ def mock_dependencies(mocker, monkeypatch):
         f"{_MODULE}.FlowMatchEulerDiscreteScheduler.from_pretrained",
         lambda *a, **k: mock_scheduler,
     )
-    monkeypatch.setattr(
+    mllm_loader = mocker.patch(
         f"{_MODULE}.Qwen3VLForConditionalGeneration.from_pretrained",
-        lambda *a, **k: mllm_wrapper,
+        return_value=mllm_wrapper,
+    )
+    mllm_config_loader = mocker.patch(
+        f"{_MODULE}.Qwen3VLConfig.from_pretrained",
+        return_value=Qwen3VLConfig(),
+    )
+    mllm_builder = mocker.patch(
+        f"{_MODULE}.create_transformers_model_with_vllm_linears",
+        return_value=nn.Linear(1, 1),
     )
     monkeypatch.setattr(
         f"{_MODULE}.Qwen3VLProcessor.from_pretrained",
@@ -74,12 +84,17 @@ def mock_dependencies(mocker, monkeypatch):
     mock_transformer_cls.return_value = mock_transformer_instance
     monkeypatch.setattr(f"{_MODULE}.BooguImageTransformer2DModel", mock_transformer_cls)
 
-    # Treat the dummy model id as a local path: skips hub prefetch.
-    mocker.patch("os.path.exists", return_value=True)
+    # Treat only the dummy model id as local. Other filesystem checks (for
+    # example lazy imports in the quantization registry) must remain real.
+    path_exists = os.path.exists
+    mocker.patch("os.path.exists", side_effect=lambda path: str(path).startswith("dummy-boogu") or path_exists(path))
 
     return {
         "inner_encoder": inner_encoder,
         "mllm_wrapper": mllm_wrapper,
+        "mllm_loader": mllm_loader,
+        "mllm_config_loader": mllm_config_loader,
+        "mllm_builder": mllm_builder,
         "processor": mock_processor,
         "vae": mock_vae,
         "scheduler": mock_scheduler,
@@ -128,6 +143,9 @@ def test_constructor_wires_components(boogu_pipeline, mock_dependencies):
     assert boogu_pipeline.vae_scale_factor == 8
     assert boogu_pipeline.default_sample_size == 128
     assert hasattr(boogu_pipeline, "load_weights")
+    assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
+    mock_dependencies["mllm_config_loader"].assert_not_called()
+    mock_dependencies["mllm_builder"].assert_not_called()
 
 
 def test_constructor_strips_mllm_lm_head(boogu_pipeline, mock_dependencies):
@@ -188,7 +206,7 @@ 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):
+def test_constructor_rejects_unsupported_mllm_quantization(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
@@ -204,10 +222,143 @@ def test_constructor_routes_only_transformer_config(mock_dependencies, mocker):
             {"transformer": transformer_config, "mllm": encoder_config, "vae": None}
         ),
     )
+    with pytest.raises(ValueError, match="Boogu MLLM only supports FP8 quantization"):
+        BooguImagePipeline(od_config=od_config)
+    mock_dependencies["mllm_loader"].assert_not_called()
+    mock_dependencies["mllm_config_loader"].assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("quantization_config", "quantize_mllm"),
+    [
+        pytest.param("fp8", True, id="global-fp8"),
+        pytest.param({"mllm": {"method": "fp8"}, "transformer": {"method": "fp8"}}, True, id="mllm-and-dit-fp8"),
+        pytest.param({"mllm": None, "transformer": "fp8"}, False, id="dit-only-fp8"),
+    ],
+)
+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
+
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        revision="test-revision",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config=quantization_config,
+    )
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    if quantize_mllm:
+        mock_dependencies["mllm_loader"].assert_not_called()
+        mock_dependencies["mllm_config_loader"].assert_called_once_with(
+            "dummy-boogu", subfolder="mllm", local_files_only=True, revision="test-revision"
+        )
+        builder = mock_dependencies["mllm_builder"]
+        builder.assert_called_once()
+        auto_cls, hf_config, online_quant_config = builder.call_args.args
+        assert auto_cls is AutoModel
+        assert hf_config is mock_dependencies["mllm_config_loader"].return_value
+        assert isinstance(online_quant_config, OnlineQuantizationConfig)
+        assert online_quant_config.args.linear.weight == kFp8Static128BlockSym
+        assert builder.call_args.kwargs == {
+            "dtype": torch.bfloat16,
+            "device": pipeline._execution_device,
+            "prefix": "mllm",
+            "skip_modules": ("mllm.visual",),
+        }
+        assert pipeline.mllm is builder.return_value
+        assert not pipeline.mllm.training
+        assert all(not parameter.requires_grad for parameter in pipeline.mllm.parameters())
+        transformer_source, mllm_source = pipeline.weights_sources
+        assert transformer_source.subfolder == "transformer"
+        assert mllm_source.model_or_path == "dummy-boogu"
+        assert mllm_source.subfolder == "mllm"
+        assert mllm_source.prefix == "mllm."
+        assert mllm_source.revision == "test-revision"
+    else:
+        assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
+        assert pipeline.mllm is mock_dependencies["inner_encoder"]
+        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)
+
+
+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
+
+    quant_config = Fp8Config(
+        ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config=quant_config,
+    )
     BooguImagePipeline(od_config=od_config)
-    kwargs = mock_dependencies["transformer_cls"].call_args.kwargs
-    assert kwargs["quant_config"] is transformer_config
-    assert kwargs["prefix"] == "transformer"
+
+    online_quant_config = mock_dependencies["mllm_builder"].call_args.args[2]
+    assert online_quant_config.ignored_layers == quant_config.ignored_layers
+
+
+def test_constructor_preserves_serialized_mllm_quantization(mock_dependencies):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    mock_dependencies["mllm_config_loader"].return_value = Qwen3VLConfig(
+        quantization_config={"quant_method": "fp8", "modules_to_not_convert": ["model.visual"]}
+    )
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu-fp8",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.bfloat16,
+        quantization_config="fp8",
+    )
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    # No override: HF reads the serialized checkpoint's scales and skip list.
+    assert "quantization_config" not in mock_dependencies["mllm_loader"].call_args.kwargs
+    assert pipeline.mllm is mock_dependencies["inner_encoder"]
+    mock_dependencies["mllm_builder"].assert_not_called()
+    assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
+
+
+def test_load_weights_maps_mllm_and_preserves_transformer():
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    pipeline = object.__new__(BooguImagePipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.mllm = nn.ModuleDict(
+        {
+            "language_model": nn.Linear(2, 2, bias=False),
+            "visual": nn.Linear(2, 2, bias=False),
+        }
+    )
+    pipeline.transformer = nn.Linear(2, 2, bias=False)
+    language_weight = torch.full((2, 2), 1.0)
+    visual_weight = torch.full((2, 2), 2.0)
+    transformer_weight = torch.full((2, 2), 3.0)
+
+    loaded = pipeline.load_weights(
+        [
+            ("mllm.lm_head.weight", torch.full((2, 2), 9.0)),
+            ("mllm.model.language_model.weight", language_weight),
+            ("mllm.model.visual.weight", visual_weight),
+            ("transformer.weight", transformer_weight),
+        ]
+    )
+
+    assert loaded == {"mllm.language_model.weight", "mllm.visual.weight", "transformer.weight"}
+    torch.testing.assert_close(pipeline.mllm["language_model"].weight, language_weight)
+    torch.testing.assert_close(pipeline.mllm["visual"].weight, visual_weight)
+    torch.testing.assert_close(pipeline.transformer.weight, transformer_weight)
 
 
 @pytest.mark.parametrize(

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -34,9 +34,13 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from torch.nn.attention import SDPBackend, sdpa_kernel
-from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor
+from transformers import AutoModel, Qwen3VLConfig, Qwen3VLForConditionalGeneration, Qwen3VLModel, Qwen3VLProcessor
+from vllm.config.quantization import QuantizationConfigArgs
 from vllm.logger import init_logger
-from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
@@ -54,6 +58,7 @@ from vllm_omni.diffusion.models.boogu_image.scheduling_flow_match_euler_discrete
 )
 from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.models.utils import create_transformers_model_with_vllm_linears
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -243,6 +248,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         self.od_config = od_config
         self._raise_unsupported_features()
         transformer_quant_config = resolve_component_quant_config(od_config.quantization_config, "transformer")
+        mllm_quant_config = resolve_component_quant_config(od_config.quantization_config, "mllm")
         self.weights_sources = [
             DiffusersPipelineLoader.ComponentSource(
                 model_or_path=od_config.model,
@@ -273,22 +279,7 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
             local_files_only=local_files_only,
             revision=od_config.revision,
         )
-
-        mllm = from_pretrained_with_prefetch(
-            Qwen3VLForConditionalGeneration.from_pretrained,
-            model,
-            subfolder="mllm",
-            prefetch_list=boogu_subfolders,
-            local_files_only=local_files_only,
-            torch_dtype=od_config.dtype,
-            revision=od_config.revision,
-        )
-        # Upstream reuses the full VLM as an optional instruction rewriter and
-        # encodes with its inner model (no ``lm_head``); the rewriter is not
-        # ported, so keep only the inner ``Qwen3VLModel`` as the encoder.
-        if hasattr(mllm, "lm_head"):
-            mllm = mllm.model
-        self.mllm = mllm.to(self._execution_device)
+        self.mllm = self._load_mllm(model, local_files_only, mllm_quant_config, boogu_subfolders)
 
         self.processor = Qwen3VLProcessor.from_pretrained(
             model,
@@ -323,6 +314,68 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         self.SYSTEM_PROMPT_4_TI2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
         self.SYSTEM_PROMPT_4_I2I = SYSTEM_PROMPT_4_TI2I_UNIFIED
 
+    def _load_mllm(
+        self,
+        model_path: str,
+        local_files_only: bool,
+        quant_config: QuantizationConfig | None,
+        prefetch_list: list[str],
+    ) -> Qwen3VLModel:
+        """Load MLLM through HF, or prepare it for deferred online FP8 loading."""
+        config = None
+        if quant_config is not None:
+            if not isinstance(quant_config, Fp8Config):
+                raise ValueError(
+                    "Boogu MLLM only supports FP8 quantization. Set mllm to null to disable online quantization."
+                )
+            config = Qwen3VLConfig.from_pretrained(
+                model_path, subfolder="mllm", local_files_only=local_files_only, revision=self.od_config.revision
+            )
+
+        # With no MLLM quantization override, or with checkpoint quantization
+        # metadata, preserve HF model initialization and weight loading.
+        if config is None or getattr(config, "quantization_config", None):
+            mllm = from_pretrained_with_prefetch(
+                Qwen3VLForConditionalGeneration.from_pretrained,
+                model_path,
+                subfolder="mllm",
+                prefetch_list=prefetch_list,
+                local_files_only=local_files_only,
+                torch_dtype=self.od_config.dtype,
+                revision=self.od_config.revision,
+            )
+            # Upstream reuses the full VLM as an optional instruction rewriter and
+            # encodes with its inner model (no ``lm_head``); the rewriter is not
+            # ported, so keep only the inner ``Qwen3VLModel`` as the encoder.
+            if hasattr(mllm, "lm_head"):
+                mllm = mllm.model
+            return mllm.to(self._execution_device)
+
+        online_quant_config = OnlineQuantizationConfig(
+            QuantizationConfigArgs(linear="fp8_per_block", ignore=quant_config.ignored_layers)
+        )
+
+        mllm = create_transformers_model_with_vllm_linears(
+            AutoModel,
+            config,
+            online_quant_config,
+            dtype=self.od_config.dtype,
+            device=self._execution_device,
+            prefix="mllm",
+            skip_modules=("mllm.visual",),
+        )
+
+        self.weights_sources.append(
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=model_path,
+                subfolder="mllm",
+                revision=self.od_config.revision,
+                prefix="mllm.",
+            )
+        )
+
+        return mllm.requires_grad_(False).eval()
+
     def _raise_unsupported_features(self) -> None:
         """Reject execution modes that do not have Boogu-specific support."""
         parallel_config = self.od_config.parallel_config
@@ -339,7 +392,8 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
-        return loader.load_weights(weights)
+        mapper = WeightsMapper(orig_to_new_prefix={"mllm.lm_head.": None, "mllm.model.": "mllm."})
+        return loader.load_weights(weights, mapper=mapper)
 
     # ------------------------------------------------------------------
     # Prompt encoding (upstream ``encode_instruction``, t2i path)

vllm_omni/diffusion/models/utils.py

@@ -5,12 +5,14 @@ from __future__ import annotations
 
 import json
 import os
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Literal
 
 import torch
 import torch.nn as nn
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
+    LinearBase,
     ReplicatedLinear,
     RowParallelLinear,
 )
@@ -124,6 +126,28 @@ def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig):
     _recursive_replace(model, prefix="")
 
 
+def recursive_replace_linear_with_quantization_config(
+    model: nn.Module,
+    quant_config: QuantizationConfig | None,
+    prefix: str = "",
+    skip_modules: Sequence[str] = (),
+) -> None:
+    for child_name, child_module in model.named_children():
+        qual_name = maybe_prefix(prefix, child_name)
+        if any(qual_name == path or qual_name.startswith(path + ".") for path in skip_modules):
+            continue
+        if isinstance(child_module, nn.Linear):
+            setattr(
+                model,
+                child_name,
+                replace_linear_class(child_module, quant_config=quant_config, prefix=qual_name),
+            )
+        else:
+            recursive_replace_linear_with_quantization_config(
+                child_module, quant_config, prefix=qual_name, skip_modules=skip_modules
+            )
+
+
 def init_parameters(
     module: nn.Module,
     dtype: torch.dtype | None,
@@ -144,6 +168,24 @@ def init_parameters(
         init_parameters(child, dtype, device)
 
 
+def init_parameters_preserving_vllm_linear(
+    module: nn.Module,
+    dtype: torch.dtype | None,
+    device: torch.device | None = None,
+) -> None:
+    if isinstance(module, LinearBase):
+        return
+    for name, param in module.named_parameters(recurse=False):
+        if param.is_meta:
+            setattr(
+                module,
+                name,
+                nn.Parameter(torch.empty_like(param, dtype=dtype, device=device), requires_grad=param.requires_grad),
+            )
+    for child in module.children():
+        init_parameters_preserving_vllm_linear(child, dtype, device)
+
+
 def create_transformers_model(
     auto_cls: _BaseAutoModelClass,
     od_config: OmniDiffusionConfig,
@@ -161,6 +203,23 @@ def create_transformers_model(
     return model
 
 
+def create_transformers_model_with_vllm_linears(
+    auto_cls: _BaseAutoModelClass,
+    hf_config: PretrainedConfig,
+    quant_config: QuantizationConfig | None,
+    dtype: torch.dtype,
+    device: torch.device,
+    prefix: str = "",
+    skip_modules: Sequence[str] = (),
+) -> PreTrainedModel:
+    """Build a Transformers model with vLLM linears."""
+    with init_on_device_without_buffers("meta"):
+        model = auto_cls.from_config(hf_config, dtype=dtype)
+    recursive_replace_linear_with_quantization_config(model, quant_config, prefix=prefix, skip_modules=skip_modules)
+    init_parameters_preserving_vllm_linear(model, dtype=dtype, device=device)
+    return model
+
+
 def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:
     """Load a JSON config file from a local path or HuggingFace Hub repo."""
     if local_files_only: