（Phase 1）Add ModelOpt FP8 auto-detect support for diffusion checkpoints #2709 (#2913)

Signed-off-by: roG0d <rodgarcas98@gmail.com>
Signed-off-by: roG0d <baonudesifeizhai@gmail.com>
Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com>
Co-authored-by: roG0d <rodgarcas98@gmail.com>

docs/.nav.yml

@@ -78,6 +78,7 @@ nav:
       - Online Quantization: user_guide/quantization/online.md
       - FP8 W8A8: user_guide/quantization/fp8.md
       - Int8 W8A8: user_guide/quantization/int8.md
+      - ModelOpt: user_guide/quantization/modelopt.md
       - GGUF: user_guide/quantization/gguf.md
       - AutoRound: user_guide/quantization/autoround.md
       - msModelSlim: user_guide/quantization/msmodelslim.md

docs/user_guide/quantization/fp8.md

@@ -2,11 +2,11 @@
 
 ## Overview
 
-FP8 quantization converts BF16/FP16 weights to FP8 at model load time, or loads
-a checkpoint whose target stage already declares an FP8 quantization config.
-Online activation scaling is the default and does not require calibration.
-Static activation scaling is supported when calibrated scale information is
-available.
+FP8 quantization converts BF16/FP16 weights to FP8 at model load time. Online
+activation scaling is the default and does not require calibration. Static
+activation scaling is supported when calibrated scale information is available.
+For ModelOpt-produced pre-quantized checkpoints, see
+[ModelOpt Quantization](modelopt.md).
 
 Some architectures can quantize all linear layers. Others have
 quality-sensitive layers that should stay in BF16 through `ignored_layers`.
@@ -46,7 +46,7 @@ guide. FP8 on Ampere may use a weight-only path where available.
 
 | Model | Scope | Format | Status |
 |-------|-------|--------|--------|
-| Qwen3-Omni | Thinker language-model stage | ModelOpt `quant_algo=FP8` | Tested for thinker memory reduction |
+| Qwen3-Omni | Thinker language-model stage | [ModelOpt](modelopt.md) `quant_algo=FP8` | Tested for thinker memory reduction |
 | Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
 
 Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a

docs/user_guide/quantization/modelopt.md

@@ -0,0 +1,138 @@
+# ModelOpt Quantization
+
+## Overview
+
+ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The
+quantized weights and scale tensors are generated before serving, so inference
+does not run online calibration or convert a BF16 checkpoint at startup.
+
+vLLM-Omni currently validates the ModelOpt FP8 checkpoint path for diffusion
+transformers. The loader auto-detects supported ModelOpt FP8 checkpoint configs
+and keeps non-transformer components, such as the tokenizer, scheduler, text
+encoder, and VAE, on the base checkpoint unless a model-specific guide says
+otherwise.
+
+!!! note
+    `--force-cutlass-fp8` is an explicit runtime override for diffusion
+    checkpoints that already carry a supported ModelOpt FP8 config. It does not
+    quantize BF16 checkpoints and it does not apply to online `--quantization
+    fp8`. The flag only takes effect for ModelOpt FP8 diffusion stages on CUDA
+    SM89+ devices; other platforms and non-ModelOpt FP8 paths fall back to the
+    normal vLLM kernel selection.
+
+## Supported ModelOpt Checkpoint Formats
+
+vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The
+checkpoint config must identify ModelOpt as the quantization method or producer,
+and the quantization algorithm must be one of the validated FP8 algorithms.
+
+| Checkpoint field | Supported value |
+|------------------|-----------------|
+| `method` / `quant_method` | `modelopt` |
+| `producer.name` | `modelopt` |
+| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` |
+
+Other ModelOpt algorithms, such as NVFP4, are not enabled by this diffusion
+FP8 path until they have separate model and quality validation.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ⭕ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
+The optional CUTLASS FP8 runtime override requires CUDA SM89+.
+
+## Model Type Support
+
+### Diffusion Model
+
+| Model | HF checkpoint | Scope | Status |
+|-------|---------------|-------|--------|
+| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints |
+| Wan2.2 | Not available | Diffusion transformer | Not validated |
+
+### Multi-Stage Omni/TTS Model
+
+| Model | Scope | Status |
+|-------|-------|--------|
+| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path |
+| Qwen3-TTS | TTS language-model stage | Not validated |
+
+Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless
+a model-specific guide documents otherwise.
+
+### Multi-Stage Diffusion Model
+
+ModelOpt checkpoints must be routed to the stage whose checkpoint contains the
+ModelOpt `quantization_config`. BAGEL and GLM-Image are not listed as validated
+ModelOpt targets yet.
+
+## Configuration
+
+For pre-quantized ModelOpt FP8 checkpoints, no `--quantization fp8` flag is
+needed. The checkpoint config selects the ModelOpt path.
+
+Online serving:
+
+```bash
+vllm serve <modelopt-fp8-checkpoint> \
+  --omni \
+  --tensor-parallel-size <N> \
+  --force-cutlass-fp8
+```
+
+Offline inference:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model <modelopt-fp8-checkpoint> \
+  --tensor-parallel-size <N> \
+  --prompt "a red ceramic teapot on a wooden table" \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 20 \
+  --seed 42 \
+  --output outputs/modelopt_fp8.png
+```
+
+Python API:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="<modelopt-fp8-checkpoint>",
+    tensor_parallel_size=2,
+    force_cutlass_fp8=True,
+)
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ |
+
+## Validation and Notes
+
+1. Compare the ModelOpt FP8 checkpoint against the BF16 baseline with the same
+   prompt, resolution, seed, and inference steps.
+2. Use `tests/diffusion/quantization/test_quantization_quality.py` with
+   `VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model
+   paths.
+3. Report LPIPS, PSNR, MAE, throughput, latency, and peak memory when adding a
+   new validated ModelOpt diffusion checkpoint.
+4. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
+   ModelOpt path only when the checkpoint already contains ModelOpt FP8 weights
+   and scales.

docs/user_guide/quantization/overview.md

@@ -10,18 +10,18 @@ type has a different quantization scope.
 | Mode | Guide | Description | Methods |
 |------|-------|-------------|---------|
 | Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8 |
-| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | GGUF, AutoRound, msModelSlim, serialized Int8 |
+| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8 |
 
 ## Hardware Support
 
-| Device | FP8 W8A8 | Int8 W8A8 | GGUF | AutoRound | msModelSlim |
-|--------|----------|-----------|------|-----------|-------------|
-| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ❌ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ❌ |
-| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ✅ | ❌ |
-| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
-| Intel XPU | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
-| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ |
+| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | GGUF | AutoRound | msModelSlim |
+|--------|----------|-----------|----------|------|-----------|-------------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
+| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
+| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
 guide. FP8 on Ampere may use a weight-only path where available.
@@ -39,6 +39,7 @@ otherwise.
 |--------|-------|------|----------------|--------|
 | FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
 | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
+| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
 | GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
 | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
@@ -52,7 +53,7 @@ in BF16 unless the model guide explicitly adds support.
 
 | Method | Guide | Scope | Example models | Status |
 |--------|-------|-------|----------------|--------|
-| FP8 | [FP8](fp8.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
+| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
 | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
 | GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
 | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
@@ -67,6 +68,7 @@ attached to the intended stage rather than applied globally.
 |--------|-------|-------|----------------|--------|
 | FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
 | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
+| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
 | GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
 | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
@@ -94,7 +96,7 @@ config = build_quant_config({
 
 | Component | Default quantized? | Notes |
 |-----------|--------------------|-------|
-| Diffusion transformer | Yes | Primary target for FP8, Int8, GGUF, AutoRound, and msModelSlim |
+| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim |
 | Text encoder | No | Keep BF16 unless a method-specific guide documents support |
 | VAE | No | Keep BF16; storage-only paths are method-specific |
 | Scheduler/tokenizer | No | Loaded from the base model repository |

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
+    ModelOptFp8CheckpointAdapter,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _PackedModelOptModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.transformer = nn.Module()
+        self.transformer.block = nn.Module()
+        self.transformer.block.to_qkv = nn.Linear(2, 2, bias=False)
+
+
+class _QuantizedPackedModelOptModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.transformer = nn.Module()
+        self.transformer.block = nn.Module()
+        self.transformer.block.to_qkv = nn.Module()
+        self.transformer.block.to_qkv.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False),
+        )
+        self.transformer.block.to_qkv.register_parameter(
+            "weight_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+        self.transformer.block.to_qkv.register_parameter(
+            "input_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+
+
+def _make_source() -> SimpleNamespace:
+    return SimpleNamespace(
+        subfolder="transformer",
+        prefix="transformer.",
+    )
+
+
+def test_modelopt_adapter_dequantizes_fp8_weight_for_full_precision_target():
+    model = _PackedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    scale = torch.tensor([0.5], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.block.to_q.weight_scale", scale),
+                    ("transformer.block.to_q.input_scale", torch.tensor([1.0])),
+                    ("transformer.block.to_q.weight", fp8_weight),
+                ]
+            )
+        )
+    )
+
+    assert [name for name, _ in adapted] == ["transformer.block.to_q.weight"]
+    assert adapted[0][1].dtype == model.transformer.block.to_qkv.weight.dtype
+    assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale)
+
+
+def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target():
+    model = _QuantizedPackedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    scale = torch.tensor([0.5], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.block.to_q.weight_scale", scale),
+                    ("transformer.block.to_q.input_scale", torch.tensor([1.0])),
+                ]
+            )
+        )
+    )
+
+    assert [name for name, _ in adapted] == [
+        "transformer.block.to_q.weight_scale",
+        "transformer.block.to_q.input_scale",
+    ]

tests/diffusion/quantization/test_fp8_config.py

@@ -54,6 +54,25 @@ def test_build_quant_config_dict_not_mutated():
     assert original == copy
 
 
+def test_build_quant_config_modelopt_fp8_config_json():
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config(
+        {
+            "quant_method": "modelopt",
+            "quant_algo": "FP8",
+            "ignore": ["proj_out"],
+            "producer": {"name": "modelopt"},
+        }
+    )
+
+    assert isinstance(config, ModelOptFp8Config)
+    assert config.get_name() == "modelopt"
+    assert config.is_checkpoint_fp8_serialized
+
+
 def test_build_quant_config_per_component():
     from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
 
@@ -91,7 +110,7 @@ def test_flat_dict_not_misdetected_as_per_component():
     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"):
+    with pytest.raises(ValueError, match="must have a 'method' or 'quant_method' key"):
         build_quant_config({"activation_scheme": "static"})
 
 
@@ -194,10 +213,30 @@ def test_integration_per_component():
     assert config.quantization_config.component_configs["vae"] is None
 
 
+def test_transformer_config_auto_detects_modelopt_fp8():
+    from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config
+
+    from vllm_omni.diffusion.data import TransformerConfig
+
+    config = TransformerConfig.from_dict(
+        {
+            "_class_name": "FluxTransformer2DModel",
+            "quantization_config": {
+                "quant_method": "modelopt",
+                "quant_algo": "FP8",
+                "ignore": ["proj_out"],
+            },
+        }
+    )
+
+    assert isinstance(config.quant_config, ModelOptFp8Config)
+    assert config.quant_method == "modelopt"
+
+
 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"]:
+    for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]:
         assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
 
 
@@ -243,14 +282,10 @@ def test_per_component_routing_with_default():
     assert resolved.get_name() == "fp8"
 
 
-@pytest.mark.parametrize(
-    "quant_algo",
-    ["FP8", "NVFP4"],
-    ids=["modelopt_fp8", "modelopt_nvfp4"],
-)
+@pytest.mark.parametrize("quant_algo", ["FP8"], ids=["modelopt_fp8"])
 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."""
+    thinker_config.text_config for verified modelopt FP8 checkpoints."""
     from types import SimpleNamespace
 
     from vllm_omni.config.model import OmniModelArchConfigConvertor

tests/diffusion/quantization/test_quantization_quality.py

@@ -1,8 +1,7 @@
 # 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.
+Quantization quality gate for diffusion models.
 
 Developers adding a new quantization method should:
 1. Add their method + model to QUALITY_CONFIGS below
@@ -11,7 +10,7 @@ Developers adding a new quantization method should:
 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.
+same seed, computes similarity metrics, and fails if LPIPS exceeds the threshold.
 
 Requirements:
     pip install lpips
@@ -21,12 +20,22 @@ Example — run only FP8 tests:
 
 Example — run a specific model:
     pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
+
+Example — validate a local BF16 baseline against a local pre-quantized checkpoint:
+    export VLLM_OMNI_QUALITY_CONFIGS=/tmp/modelopt_quality_cases.json
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "qwen_image_2512"
+
+Optional artifact dump:
+    export VLLM_OMNI_QUALITY_OUTPUT_DIR=/tmp/modelopt_quality_outputs
 """
 
 from __future__ import annotations
 
 import gc
+import json
+import os
 from dataclasses import dataclass
+from pathlib import Path
 
 import numpy as np
 import pytest
@@ -44,17 +53,49 @@ 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
+    model: str | None = None  # HF model name
+    quantization: str | None = None  # quantization method, e.g. "fp8"
+    baseline_model: str | None = None  # explicit BF16/local baseline path
+    quantized_model: str | None = None  # explicit quantized/local model path
     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
+    negative_prompt: str = ""
+    guidance_scale: float | None = None
+
+    def baseline_ref(self) -> str:
+        return self.baseline_model or self.model or ""
+
+    def quantized_ref(self) -> str:
+        if self.quantized_model is not None:
+            return self.quantized_model
+        return self.model or ""
+
+    def quantization_ref(self) -> str | None:
+        if self.quantized_model is not None:
+            return None
+        return self.quantization
+
+    def validate(self) -> None:
+        uses_explicit_models = self.baseline_model is not None or self.quantized_model is not None
+        uses_model_plus_method = self.model is not None or self.quantization is not None
+
+        if uses_explicit_models and uses_model_plus_method:
+            raise ValueError(f"{self.id}: explicit baseline/quantized paths cannot be mixed with model/quantization")
+
+        if uses_explicit_models:
+            if self.baseline_model is None or self.quantized_model is None:
+                raise ValueError(f"{self.id}: baseline_model and quantized_model must be provided together")
+            return
+
+        if self.model is None or self.quantization is None:
+            raise ValueError(f"{self.id}: expected either model+quantization or baseline_model+quantized_model")
 
 
 # Add new quantization methods / models here.
@@ -91,6 +132,53 @@ QUALITY_CONFIGS = [
 ]
 
 
+def _load_extra_quality_configs() -> list[QualityTestConfig]:
+    config_path = os.environ.get("VLLM_OMNI_QUALITY_CONFIGS")
+    if not config_path:
+        return []
+
+    payload = json.loads(Path(config_path).read_text(encoding="utf-8"))
+    if not isinstance(payload, list):
+        raise ValueError("VLLM_OMNI_QUALITY_CONFIGS must point to a JSON list")
+
+    configs: list[QualityTestConfig] = []
+    for item in payload:
+        if not isinstance(item, dict):
+            raise ValueError("Each quality config entry must be a JSON object")
+        config = QualityTestConfig(**item)
+        config.validate()
+        configs.append(config)
+    return configs
+
+
+def _all_quality_configs() -> list[QualityTestConfig]:
+    configs = [*QUALITY_CONFIGS, *_load_extra_quality_configs()]
+    for config in configs:
+        config.validate()
+    return configs
+
+
+def _output_path(output_dir: Path, config: QualityTestConfig, label: str, suffix: str) -> Path:
+    safe_id = config.id.replace("/", "_")
+    return output_dir / f"{safe_id}_{label}{suffix}"
+
+
+def _maybe_save_output(output_dir: Path | None, config: QualityTestConfig, label: str, output) -> None:
+    if output_dir is None:
+        return
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    if config.task == "t2i":
+        output.save(_output_path(output_dir, config, label, ".png"))
+        return
+
+    if isinstance(output, np.ndarray):
+        np.save(_output_path(output_dir, config, label, ".npy"), output)
+        return
+
+    raise TypeError(f"Unsupported output type for saving: {type(output)!r}")
+
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -107,12 +195,13 @@ def _generate_image(omni, config: QualityTestConfig):
     torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
-        {"prompt": config.prompt},
+        {"prompt": config.prompt, "negative_prompt": config.negative_prompt},
         OmniDiffusionSamplingParams(
             height=config.height,
             width=config.width,
             generator=generator,
             num_inference_steps=config.num_inference_steps,
+            guidance_scale=config.guidance_scale,
         ),
     )
 
@@ -138,13 +227,14 @@ def _generate_video(omni, config: QualityTestConfig):
     torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
-        {"prompt": config.prompt, "negative_prompt": ""},
+        {"prompt": config.prompt, "negative_prompt": config.negative_prompt},
         OmniDiffusionSamplingParams(
             height=config.height,
             width=config.width,
             generator=generator,
             num_inference_steps=config.num_inference_steps,
             num_frames=config.num_frames,
+            guidance_scale=config.guidance_scale,
         ),
     )
 
@@ -186,6 +276,34 @@ def _compute_lpips(baseline, quantized, task: str) -> float:
     return compute_lpips_video(baseline, quantized)
 
 
+def _to_float_array(output, task: str) -> np.ndarray:
+    if task == "t2i":
+        array = np.asarray(output.convert("RGB"), dtype=np.float32) / 255.0
+    else:
+        array = np.asarray(output, dtype=np.float32)
+        if array.max() > 1.0 or array.min() < 0.0:
+            array = np.clip(array, 0.0, 255.0) / 255.0
+        else:
+            array = np.clip(array, 0.0, 1.0)
+    return array
+
+
+def _compute_psnr_and_mae(baseline, quantized, task: str) -> tuple[float, float]:
+    baseline_array = _to_float_array(baseline, task)
+    quantized_array = _to_float_array(quantized, task)
+    if baseline_array.shape != quantized_array.shape:
+        raise ValueError(
+            "Output shapes do not match for metric computation: "
+            f"baseline={baseline_array.shape}, quantized={quantized_array.shape}"
+        )
+
+    diff = baseline_array - quantized_array
+    mae = float(np.mean(np.abs(diff)))
+    mse = float(np.mean(np.square(diff)))
+    psnr = float("inf") if mse == 0.0 else float(20.0 * np.log10(1.0 / np.sqrt(mse)))
+    return psnr, mae
+
+
 def _unload(omni):
     del omni
     gc.collect()
@@ -199,6 +317,7 @@ def _unload(omni):
 # ---------------------------------------------------------------------------
 
 _marks = hardware_marks(res={"cuda": "H100"})
+_OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
 
 
 def _quality_param(c: QualityTestConfig):
@@ -214,7 +333,7 @@ def _quality_param(c: QualityTestConfig):
 @pytest.mark.diffusion
 @pytest.mark.parametrize(
     "config",
-    [_quality_param(c) for c in QUALITY_CONFIGS],
+    [_quality_param(c) for c in _all_quality_configs()],
 )
 def test_quantization_quality(config: QualityTestConfig):
     """Validate that quantized output stays within LPIPS threshold of BF16."""
@@ -223,31 +342,44 @@ def test_quantization_quality(config: QualityTestConfig):
     generate_fn = _generate_video if config.task == "t2v" else _generate_image
 
     # --- BF16 baseline ---
-    omni_bl = Omni(model=config.model)
+    omni_bl = Omni(model=config.baseline_ref())
     baseline_out, bl_mem = generate_fn(omni_bl, config)
     _unload(omni_bl)
+    _maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out)
 
     # --- Quantized ---
-    omni_qt = Omni(model=config.model, quantization_config=config.quantization)
+    quantization = config.quantization_ref()
+    if quantization is None:
+        omni_qt = Omni(model=config.quantized_ref())
+    else:
+        omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization)
     quant_out, qt_mem = generate_fn(omni_qt, config)
     _unload(omni_qt)
+    _maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out)
 
-    # --- LPIPS ---
+    # --- Similarity metrics ---
     lpips_score = _compute_lpips(baseline_out, quant_out, config.task)
+    psnr_score, mae_score = _compute_psnr_and_mae(baseline_out, quant_out, config.task)
+    assert lpips_score <= config.max_lpips, (
+        f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} "
+        f"for {config.quantization_ref() or 'pre-quantized checkpoint'} on {config.quantized_ref()}"
+    )
 
     # --- 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"  Baseline:      {config.baseline_ref()}")
+    print(f"  Quantized:     {config.quantized_ref()}")
+    print(f"  Method:        {config.quantization_ref() or 'pre-quantized checkpoint'}")
     print(f"  LPIPS:         {lpips_score:.4f}  (threshold: {config.max_lpips})")
+    print(f"  PSNR:          {psnr_score:.4f} dB  (higher is better)")
+    print(f"  MAE:           {mae_score:.6f}  (lower is better)")
     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}"
-    )
+    assert np.isfinite(psnr_score) or np.isinf(psnr_score), f"PSNR is invalid for {config.id}: {psnr_score}"
+    assert np.isfinite(mae_score), f"MAE is not finite for {config.id}: {mae_score}"

tests/diffusion/test_diffusion_worker.py

@@ -14,7 +14,7 @@ import pytest
 import torch
 from pytest_mock import MockerFixture
 
-from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker, _make_diffusion_vllm_model_config
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 
@@ -78,6 +78,34 @@ class TestDiffusionWorkerLoadWeights:
         assert result == set()
 
 
+def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods():
+    from types import SimpleNamespace
+
+    from vllm_omni.quantization import build_quant_config
+
+    od_config = SimpleNamespace(
+        model="dummy",
+        dtype=torch.bfloat16,
+        quantization_config=build_quant_config(
+            {
+                "quant_method": "modelopt",
+                "quant_algo": "FP8",
+                "ignore": [],
+            }
+        ),
+        tf_model_config=SimpleNamespace(),
+        enforce_eager=True,
+        is_moe=False,
+    )
+
+    model_config = _make_diffusion_vllm_model_config(od_config)
+
+    assert model_config.dtype is torch.bfloat16
+    assert model_config.quantization == "modelopt"
+    assert model_config.quantization_config is od_config.quantization_config
+    assert model_config.is_quantized()
+
+
 class TestDiffusionWorkerSleep:
     """Test DiffusionWorker.sleep method."""
 

vllm_omni/diffusion/data.py

@@ -208,10 +208,11 @@ class TransformerConfig:
         quant_method: str | None = None
         quant_config: QuantizationConfig | None = None
         disk_qc = params.get("quantization_config")
-        if isinstance(disk_qc, dict) and "quant_method" in disk_qc:
-            quant_method = disk_qc["quant_method"]
-            kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
-            quant_config = build_quant_config(quant_method, **kwargs)
+        if isinstance(disk_qc, dict):
+            raw_quant_method = disk_qc.get("quant_method", disk_qc.get("method"))
+            quant_config = build_quant_config(disk_qc)
+            if quant_config is not None:
+                quant_method = raw_quant_method if raw_quant_method is not None else quant_config.get_name()
 
         return cls(params=params, quant_method=quant_method, quant_config=quant_config)
 
@@ -513,6 +514,10 @@ class OmniDiffusionConfig:
     # str is resolved to {"method": <str>} internally.
     # Per-component: {"transformer": {"method": "fp8"}, "vae": None}
     quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
+    # Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This
+    # does not enable FP8 by itself; it only selects CUTLASS once the checkpoint
+    # has already resolved to vLLM's ModelOpt FP8 linear method.
+    force_cutlass_fp8: bool = False
 
     # Diffusion pipeline Profiling config
     enable_diffusion_pipeline_profiler: bool = False
@@ -630,14 +635,9 @@ class OmniDiffusionConfig:
 
         # Auto-detect quantization from TransformerConfig if not explicitly set.
         # This covers the case where tf_model_config is passed at construction
-        # time.  For late (post-construction) assignment, callers should use
+        # time. For late (post-construction) assignment, callers should use
         # set_tf_model_config() which propagates quant_config automatically.
-        if self.quantization_config is None and self.tf_model_config.quant_config is not None:
-            self.quantization_config = self.tf_model_config.quant_config
-            logger.info(
-                "Auto-detected quantization '%s' from model config",
-                self.tf_model_config.quant_method,
-            )
+        self._propagate_quantization_from_tf_config(self.tf_model_config)
 
         # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.
         if self.quantization_config is not None:
@@ -664,6 +664,32 @@ class OmniDiffusionConfig:
                 "valid together with diffusion_load_format=diffusers"
             )
 
+    def _propagate_quantization_from_tf_config(self, tf_config: "TransformerConfig") -> None:
+        if tf_config.quant_config is None:
+            return
+
+        is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, "is_checkpoint_fp8_serialized", False))
+        should_use_checkpoint_config = self.quantization_config is None or (
+            is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config)
+        )
+        if should_use_checkpoint_config:
+            self.quantization_config = tf_config.quant_config
+            logger.info(
+                "Auto-detected quantization '%s' from model config",
+                tf_config.quant_method,
+            )
+
+    @staticmethod
+    def _is_generic_fp8_quant_config(quant_config: object) -> bool:
+        if isinstance(quant_config, str):
+            return quant_config.lower() == "fp8"
+        if isinstance(quant_config, Mapping):
+            method = quant_config.get("method", quant_config.get("quant_method"))
+            return isinstance(method, str) and method.lower() == "fp8"
+        if hasattr(quant_config, "get_name"):
+            return quant_config.get_name() == "fp8"
+        return False
+
     def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
         """Assign `tf_model_config` and propagate quantization if detected.
 
@@ -679,12 +705,7 @@ class OmniDiffusionConfig:
                 `TransformerConfig.from_dict`.
         """
         self.tf_model_config = tf_config
-        if self.quantization_config is None and tf_config.quant_config is not None:
-            self.quantization_config = tf_config.quant_config
-            logger.info(
-                "Auto-detected quantization '%s' from model config",
-                tf_config.quant_method,
-            )
+        self._propagate_quantization_from_tf_config(tf_config)
 
     def update_multimodal_support(self) -> None:
         # Resolve serving-visible multimodal behavior from shared metadata
@@ -716,50 +737,35 @@ class OmniDiffusionConfig:
                 # Skip transformer config loading for diffusers adapter
                 # (non-DiT models don't have a separate transformer folder/config)
                 if self.diffusion_load_format == "diffusers":
-                    self.tf_model_config = TransformerConfig()
-                    try:
-                        diffusers_pipeline_cls_name = config_dict["_class_name"]
-                        self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name)
-                    except (KeyError, AttributeError) as exc:
-                        logger.warning(
-                            f"Could not find valid _class_name for diffusers pipeline in model_index.json: {exc}. "
-                            "Without knowing the underlying diffusers pipeline class, the dummy run will input only "
-                            "text prompt, which may cause errors for pipelines that require additional inputs."
-                        )
+                    self.set_tf_model_config(TransformerConfig())
                 else:
                     tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
-                    self.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
+                    self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
             else:
                 raise FileNotFoundError("model_index.json not found")
         except (AttributeError, OSError, ValueError, FileNotFoundError):
             # Skip transformer config loading for diffusers adapter
             # (non-DiT models don't have a separate transformer folder/config)
             if self.diffusion_load_format == "diffusers":
-                self.tf_model_config = TransformerConfig()
-                logger.warning(
-                    "Could not find valid model_index.json per diffusers format. "
-                    "This model is suspectedly unsupported by the diffusers backend. "
-                    "Also, without knowing the underlying diffusers pipeline class from model_index.json, "
-                    "the dummy run will input only text prompt, which may cause errors for pipelines "
-                    "that require additional inputs."
-                )
+                self.set_tf_model_config(TransformerConfig())
+                self.update_multimodal_support()
             else:
                 cfg = get_hf_file_to_dict("config.json", self.model)
                 if cfg is None:
                     raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
 
-                self.tf_model_config = TransformerConfig.from_dict(cfg)
+                self.set_tf_model_config(TransformerConfig.from_dict(cfg))
                 model_type = cfg.get("model_type")
                 architectures = cfg.get("architectures") or []
 
                 if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
                     self.model_class_name = "BagelPipeline"
-                    self.tf_model_config = TransformerConfig()
+                    self.set_tf_model_config(TransformerConfig())
                     self.update_multimodal_support()
                 elif model_type == "nextstep":
                     if self.model_class_name is None:
                         self.model_class_name = "NextStep11Pipeline"
-                    self.tf_model_config = TransformerConfig()
+                    self.set_tf_model_config(TransformerConfig())
                     self.update_multimodal_support()
                 elif model_type == "s2v":
                     if self.model_class_name is None:
@@ -794,23 +800,10 @@ class OmniDiffusionConfig:
             cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
             kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
 
-        # Falsy-value check for not-None fields (convert potential None values in YAML config to empty containers)
-        if "diffusers_load_kwargs" in kwargs and kwargs["diffusers_load_kwargs"] is None:
-            kwargs["diffusers_load_kwargs"] = {}
-        if "diffusers_call_kwargs" in kwargs and kwargs["diffusers_call_kwargs"] is None:
-            kwargs["diffusers_call_kwargs"] = {}
-
-        # Forward top-level parallel knobs (e.g. --tensor-parallel-size from CLI)
-        # into parallel_config so the diffusion engine sees them.
-        par = kwargs.get("parallel_config", {})
-        if isinstance(par, Mapping):
-            par = dict(par)
-            if par.get("tensor_parallel_size") is None:
-                par.pop("tensor_parallel_size", None)
-            tensor_parallel_size = kwargs.get("tensor_parallel_size")
-            if tensor_parallel_size is not None and "tensor_parallel_size" not in par:
-                par["tensor_parallel_size"] = tensor_parallel_size
-            kwargs["parallel_config"] = par
+        # Convert optional YAML null values to empty containers.
+        for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
+            if key in kwargs and kwargs[key] is None:
+                kwargs[key] = {}
 
         # Filter kwargs to only include valid fields
         valid_fields = {f.name for f in fields(cls)}

vllm_omni/diffusion/layers/rope.py

@@ -65,6 +65,22 @@ def apply_rotary_emb_mindiesd(
         return rotary_position_embedding(x, cos, sin, rotated_mode="rotated_half", head_first=False, fused=True)
 
 
+def _ensure_batch_dim(x: torch.Tensor) -> tuple[torch.Tensor, bool]:
+    # Upstream fused rotary kernels expect ``x`` shaped as
+    # ``[batch_size, seq_len, nheads, headdim]``. Some omni diffusion call
+    # sites pass ``[seq_len, nheads, headdim]`` instead, so normalize to 4D
+    # here before entering the fused path.
+    if x.dim() == 3:
+        return x.unsqueeze(0), True
+    return x, False
+
+
+def _restore_batch_dim(x: torch.Tensor, squeezed: bool) -> torch.Tensor:
+    if squeezed:
+        return x.squeeze(0)
+    return x
+
+
 class RotaryEmbedding(CustomOp):
     """
     rotary positional embedding.
@@ -98,12 +114,14 @@ class RotaryEmbedding(CustomOp):
             cos = cos[0]
             sin = sin[0]
 
-        return apply_rotary_emb(
+        x, squeezed = _ensure_batch_dim(x)
+        output = apply_rotary_emb(
             x,
             cos,
             sin,
             interleaved=self.interleaved,
         )
+        return _restore_batch_dim(output, squeezed)
 
     def forward_hip(
         self,
@@ -119,12 +137,14 @@ class RotaryEmbedding(CustomOp):
             cos = cos[0]
             sin = sin[0]
 
-        return self.apply_rotary_emb_flash_attn(
+        x, squeezed = _ensure_batch_dim(x)
+        output = self.apply_rotary_emb_flash_attn(
             x,
             cos,
             sin,
             interleaved=self.interleaved,
         )
+        return _restore_batch_dim(output, squeezed)
 
     def forward_npu(
         self,

vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from torch import nn
+
+from .modelopt import ModelOptFp8CheckpointAdapter
+
+
+def get_checkpoint_adapter(
+    model: nn.Module,
+    source: object,
+    quant_config: object | None,
+    use_safetensors: bool,
+) -> ModelOptFp8CheckpointAdapter | None:
+    if ModelOptFp8CheckpointAdapter.is_compatible(source, quant_config, use_safetensors):
+        return ModelOptFp8CheckpointAdapter(model, source)
+    return None
+
+
+__all__ = [
+    "ModelOptFp8CheckpointAdapter",
+    "get_checkpoint_adapter",
+]

vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass, field
+
+import torch
+from torch import nn
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.model_executor.utils import get_packed_modules_mapping
+
+logger = init_logger(__name__)
+
+MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_inv")
+DEFAULT_PACKED_MODULES_MAPPING = {
+    "to_qkv": ("to_q", "to_k", "to_v"),
+    "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
+    "w13": ("w1", "w3"),
+}
+FP8_DTYPES = tuple(
+    dtype
+    for dtype in (
+        getattr(torch, "float8_e4m3fn", None),
+        getattr(torch, "float8_e5m2", None),
+        getattr(torch, "float8_e4m3fnuz", None),
+        getattr(torch, "float8_e5m2fnuz", None),
+    )
+    if dtype is not None
+)
+
+
+@dataclass
+class _AdaptState:
+    scale_tensors: dict[str, torch.Tensor] = field(default_factory=dict)
+    pending_weights: dict[str, list[tuple[str, torch.Tensor, torch.dtype]]] = field(default_factory=dict)
+    skipped_scales: int = 0
+    dequantized_weights: int = 0
+
+
+class ModelOptFp8CheckpointAdapter:
+    def __init__(self, model: nn.Module, source: object):
+        self._loadable_tensors = self._get_model_loadable_tensors(model)
+        self._weights_mapper = self._get_weights_mapper(model)
+        self._source_label = getattr(source, "prefix", "") or getattr(source, "subfolder", None) or "model"
+
+    @classmethod
+    def is_compatible(
+        cls,
+        source: object,
+        quant_config: object | None,
+        use_safetensors: bool,
+    ) -> bool:
+        return use_safetensors and cls._is_transformer_source(source) and cls._is_checkpoint_quant_config(quant_config)
+
+    @staticmethod
+    def _is_transformer_source(source: object) -> bool:
+        if getattr(source, "subfolder", None) == "transformer":
+            return True
+        return str(getattr(source, "prefix", "")).startswith("transformer.")
+
+    @staticmethod
+    def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
+        return (
+            quant_config is not None
+            and hasattr(quant_config, "get_name")
+            and quant_config.get_name() == "modelopt"
+            and bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
+        )
+
+    @staticmethod
+    def _get_model_loadable_tensors(model: nn.Module) -> dict[str, torch.Tensor]:
+        loadable_tensors: dict[str, torch.Tensor] = {name: param for name, param in model.named_parameters()}
+        loadable_tensors.update({name: buffer for name, buffer in model.named_buffers()})
+        return loadable_tensors
+
+    @staticmethod
+    def _is_scale(name: str) -> bool:
+        return name.endswith(MODEL_OPT_SCALE_SUFFIXES)
+
+    @staticmethod
+    def _is_fp8_tensor(tensor: torch.Tensor) -> bool:
+        return tensor.dtype in FP8_DTYPES
+
+    @staticmethod
+    def _get_weight_scale_name(weight_name: str) -> str | None:
+        if weight_name.endswith(".weight"):
+            return weight_name[: -len(".weight")] + ".weight_scale"
+        return None
+
+    @classmethod
+    def _get_weights_mapper(cls, model: nn.Module) -> WeightsMapper:
+        mapping = {
+            packed_name: tuple(shard_names) for packed_name, shard_names in DEFAULT_PACKED_MODULES_MAPPING.items()
+        }
+        mapping.update(
+            {
+                str(packed_name): tuple(str(shard_name) for shard_name in shard_names)
+                for packed_name, shard_names in get_packed_modules_mapping(model).items()
+            }
+        )
+
+        orig_to_new_substr = {".to_out.0.": ".to_out."}
+        orig_to_new_prefix: dict[str, str] = {}
+        for packed_name, shard_names in mapping.items():
+            for shard_name in shard_names:
+                orig_to_new_substr[f".{shard_name}."] = f".{packed_name}."
+                orig_to_new_prefix[f"{shard_name}."] = f"{packed_name}."
+        return WeightsMapper(
+            orig_to_new_substr=orig_to_new_substr,
+            orig_to_new_prefix=orig_to_new_prefix,
+        )
+
+    def _resolve_target_name(self, name: str) -> str | None:
+        if name in self._loadable_tensors:
+            return name
+
+        for candidate in self._weights_mapper.apply_list([name]):
+            if candidate != name and candidate in self._loadable_tensors:
+                return candidate
+        return None
+
+    @staticmethod
+    def _reshape_weight_scale(scale: torch.Tensor, weight_shape: torch.Size) -> torch.Tensor:
+        if scale.numel() == 1:
+            return scale.reshape(())
+        if len(weight_shape) == 2 and scale.ndim == 1 and scale.shape[0] == weight_shape[0]:
+            return scale.reshape(-1, 1)
+        if tuple(scale.shape) == tuple(weight_shape):
+            return scale
+        if (
+            len(weight_shape) == 2
+            and scale.ndim == 4
+            and scale.shape[1] == 1
+            and scale.shape[3] == 1
+            and weight_shape[0] % scale.shape[0] == 0
+            and weight_shape[1] % scale.shape[2] == 0
+        ):
+            block_n = weight_shape[0] // scale.shape[0]
+            block_k = weight_shape[1] // scale.shape[2]
+            return scale.expand(scale.shape[0], block_n, scale.shape[2], block_k).reshape(weight_shape)
+        raise ValueError(f"Unsupported ModelOpt FP8 weight_scale shape {tuple(scale.shape)} for weight {weight_shape}")
+
+    def _dequantize_weight(
+        self,
+        name: str,
+        loaded_weight: torch.Tensor,
+        state: _AdaptState,
+        target_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        scale_name = self._get_weight_scale_name(name)
+        if scale_name is None or scale_name not in state.scale_tensors:
+            raise ValueError(f"Missing ModelOpt FP8 weight_scale for full-precision target weight {name!r}")
+
+        scale = state.scale_tensors[scale_name].to(dtype=target_dtype, device=loaded_weight.device)
+        scale = self._reshape_weight_scale(scale, loaded_weight.shape)
+        return loaded_weight.to(dtype=target_dtype) * scale
+
+    def _flush_pending_weights(
+        self,
+        scale_name: str,
+        state: _AdaptState,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        for weight_name, weight_tensor, target_dtype in state.pending_weights.pop(scale_name, []):
+            yield weight_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype)
+            state.dequantized_weights += 1
+
+    def _handle_scale_tensor(
+        self,
+        name: str,
+        tensor: torch.Tensor,
+        target_name: str | None,
+        state: _AdaptState,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        state.scale_tensors[name] = tensor
+        if target_name is None:
+            state.skipped_scales += 1
+        else:
+            yield name, tensor
+        yield from self._flush_pending_weights(name, state)
+
+    def _target_dtype_for_dequantization(
+        self,
+        tensor: torch.Tensor,
+        target_name: str | None,
+    ) -> torch.dtype | None:
+        if target_name is None or not self._is_fp8_tensor(tensor):
+            return None
+
+        target_dtype = self._loadable_tensors[target_name].dtype
+        if target_dtype in FP8_DTYPES:
+            return None
+        return target_dtype
+
+    def _maybe_dequantize_or_defer_weight(
+        self,
+        name: str,
+        tensor: torch.Tensor,
+        target_dtype: torch.dtype,
+        state: _AdaptState,
+    ) -> torch.Tensor | None:
+        scale_name = self._get_weight_scale_name(name)
+        if scale_name is None:
+            raise ValueError(f"Missing ModelOpt FP8 weight_scale name for weight {name!r}")
+
+        if scale_name not in state.scale_tensors:
+            state.pending_weights.setdefault(scale_name, []).append((name, tensor, target_dtype))
+            return None
+
+        state.dequantized_weights += 1
+        return self._dequantize_weight(name, tensor, state, target_dtype)
+
+    @staticmethod
+    def _check_pending_weights(state: _AdaptState) -> None:
+        if not state.pending_weights:
+            return
+
+        missing_scale_names = ", ".join(repr(name) for name in sorted(state.pending_weights))
+        raise ValueError(f"Missing ModelOpt FP8 weight_scale for full-precision target weights: {missing_scale_names}")
+
+    def _log_adaptation_summary(self, state: _AdaptState) -> None:
+        if not state.skipped_scales and not state.dequantized_weights:
+            return
+
+        logger.info_once(
+            "Adapted ModelOpt FP8 %s weights: dequantized %d full-precision weights, skipped %d scale tensors",
+            self._source_label,
+            state.dequantized_weights,
+            state.skipped_scales,
+        )
+
+    def adapt(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        state = _AdaptState()
+
+        for name, tensor in weights:
+            target_name = self._resolve_target_name(name)
+            if self._is_scale(name):
+                yield from self._handle_scale_tensor(name, tensor, target_name, state)
+                continue
+
+            target_dtype = self._target_dtype_for_dequantization(tensor, target_name)
+            if target_dtype is not None:
+                tensor = self._maybe_dequantize_or_defer_weight(name, tensor, target_dtype, state)
+                if tensor is None:
+                    continue
+            yield name, tensor
+
+        self._check_pending_weights(state)
+        self._log_adaptation_summary(state)

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -7,7 +7,7 @@ import re
 import time
 from collections.abc import Generator, Iterable
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
+from typing import cast
 
 import torch
 from huggingface_hub import hf_hub_download
@@ -32,14 +32,14 @@ from vllm.utils.torch_utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
+    get_checkpoint_adapter,
+)
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import initialize_model
 
-if TYPE_CHECKING:
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
 logger = init_logger(__name__)
 
 
@@ -177,7 +177,11 @@ class DiffusersPipelineLoader:
 
         return hf_folder, hf_weights_files, use_safetensors
 
-    def _get_weights_iterator(self, source: "ComponentSource") -> Generator[tuple[str, torch.Tensor], None, None]:
+    def _get_weights_iterator(
+        self,
+        source: "ComponentSource",
+        model: nn.Module | None = None,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
         """Get an iterator for the model weights based on the load format."""
         _, hf_weights_files, use_safetensors = self._prepare_weights(
             source.model_or_path,
@@ -213,7 +217,34 @@ class DiffusersPipelineLoader:
         if self.counter_before_loading_weights == 0.0:
             self.counter_before_loading_weights = time.perf_counter()
         # Apply the prefix.
-        return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
+        prefixed_weights_iterator = ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
+        if model is not None:
+            checkpoint_adapter = self._get_checkpoint_adapter(model, source, use_safetensors)
+            if checkpoint_adapter is not None:
+                return checkpoint_adapter.adapt(prefixed_weights_iterator)
+        return prefixed_weights_iterator
+
+    def _get_source_quant_config(self, source: "ComponentSource") -> object | None:
+        if self.od_config is None:
+            return None
+
+        quant_config = self.od_config.quantization_config
+        if hasattr(quant_config, "resolve"):
+            return quant_config.resolve(source.prefix.rstrip("."))
+        return quant_config
+
+    def _get_checkpoint_adapter(
+        self,
+        model: nn.Module,
+        source: "ComponentSource",
+        use_safetensors: bool,
+    ):
+        return get_checkpoint_adapter(
+            model=model,
+            source=source,
+            quant_config=self._get_source_quant_config(source),
+            use_safetensors=use_safetensors,
+        )
 
     def get_all_weights(
         self,
@@ -221,7 +252,7 @@ class DiffusersPipelineLoader:
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
         sources = self._get_weight_sources(model)
         for source in sources:
-            yield from self._get_weights_iterator(source)
+            yield from self._get_weights_iterator(source, model=model)
 
     def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]:
         return tuple(
@@ -269,7 +300,7 @@ class DiffusersPipelineLoader:
         """Load a model with the given configurations."""
         if load_format is None:
             load_format = "default"
-
+        self.od_config = od_config
         # CPU offload + FP8: load weights on device for FP8 quantization
         if load_device == "cpu" and od_config.quantization_config is not None:
             load_device = device.type
@@ -291,13 +322,12 @@ class DiffusersPipelineLoader:
                         model_cls = resolve_obj_by_qualname(custom_pipeline_name)
                         model = model_cls(od_config=od_config)
                     else:
-                        # 'dummy' format should not call this function at all
                         raise ValueError(f"Unknown load_format: {load_format}")
                 logger.debug("Loading weights on %s ...", load_device)
                 if load_format == "diffusers":
                     # DiffusersAdapterPipeline.load_weights() calls
-                    # DiffusionPipeline.from_pretrained() internally — it does
-                    # NOT use our native (customized) pipeline classes.
+                    # DiffusionPipeline.from_pretrained() internally; it does
+                    # not use our native customized pipeline classes.
                     cast(DiffusersAdapterPipeline, model).load_weights()
                 elif self._is_gguf_quantization(od_config):
                     self._load_weights_with_gguf(model, od_config)
@@ -449,9 +479,7 @@ class DiffusersPipelineLoader:
     def _get_model_loadable_names(self, model: nn.Module) -> set[str]:
         # Avoid model.state_dict() here because GGUF uses UninitializedParameter
         # which raises during detach(). Collect names directly.
-        names = {name for name, _ in model.named_parameters()}
-        names.update(name for name, _ in model.named_buffers())
-        return names
+        return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()}
 
     def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str:
         if os.path.isfile(gguf_model):

vllm_omni/diffusion/models/flux2/flux2_transformer.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable
 from types import SimpleNamespace
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 from diffusers.models.embeddings import (
@@ -27,6 +27,13 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+
+def _join_prefix(prefix: str, name: str) -> str:
+    return f"{prefix}.{name}" if prefix else name
+
 
 class Flux2SwiGLU(nn.Module):
     """SwiGLU activation used by Flux2."""
@@ -48,6 +55,8 @@ class Flux2FeedForward(nn.Module):
         mult: float = 3.0,
         inner_dim: int | None = None,
         bias: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         if inner_dim is None:
@@ -59,6 +68,8 @@ class Flux2FeedForward(nn.Module):
             [inner_dim, inner_dim],
             bias=bias,
             return_bias=False,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "linear_in"),
         )
         self.act_fn = Flux2SwiGLU()
         self.linear_out = RowParallelLinear(
@@ -67,6 +78,8 @@ class Flux2FeedForward(nn.Module):
             bias=bias,
             input_is_parallel=True,
             return_bias=False,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "linear_out"),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -89,6 +102,8 @@ class Flux2Attention(nn.Module):
         eps: float = 1e-5,
         out_dim: int = None,
         elementwise_affine: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.head_dim = dim_head
@@ -104,6 +119,8 @@ class Flux2Attention(nn.Module):
             head_size=self.head_dim,
             total_num_heads=self.heads,
             bias=bias,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "to_qkv"),
         )
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -119,6 +136,8 @@ class Flux2Attention(nn.Module):
                     bias=out_bias,
                     input_is_parallel=True,
                     return_bias=False,
+                    quant_config=quant_config,
+                    prefix=_join_prefix(prefix, "to_out.0"),
                 ),
                 nn.Dropout(dropout),
             ]
@@ -132,6 +151,8 @@ class Flux2Attention(nn.Module):
                 head_size=self.head_dim,
                 total_num_heads=self.heads,
                 bias=added_proj_bias,
+                quant_config=quant_config,
+                prefix=_join_prefix(prefix, "add_kv_proj"),
             )
             self.add_query_num_heads = self.add_kv_proj.num_heads
             self.add_kv_num_heads = self.add_kv_proj.num_kv_heads
@@ -141,6 +162,8 @@ class Flux2Attention(nn.Module):
                 bias=out_bias,
                 input_is_parallel=True,
                 return_bias=False,
+                quant_config=quant_config,
+                prefix=_join_prefix(prefix, "to_add_out"),
             )
 
         self.rope = RotaryEmbedding(is_neox_style=False)
@@ -238,6 +261,8 @@ class Flux2ParallelSelfAttention(nn.Module):
         elementwise_affine: bool = True,
         mlp_ratio: float = 4.0,
         mlp_mult_factor: int = 2,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.head_dim = dim_head
@@ -256,6 +281,8 @@ class Flux2ParallelSelfAttention(nn.Module):
             self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor,
             bias=bias,
             gather_output=True,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "to_qkv_mlp_proj"),
         )
         self.mlp_act_fn = Flux2SwiGLU()
 
@@ -267,6 +294,8 @@ class Flux2ParallelSelfAttention(nn.Module):
             self.out_dim,
             bias=out_bias,
             gather_output=True,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "to_out"),
         )
         self.rope = RotaryEmbedding(is_neox_style=False)
         self.attn = Attention(
@@ -329,6 +358,8 @@ class Flux2SingleTransformerBlock(nn.Module):
         mlp_ratio: float = 3.0,
         eps: float = 1e-6,
         bias: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -342,6 +373,8 @@ class Flux2SingleTransformerBlock(nn.Module):
             eps=eps,
             mlp_ratio=mlp_ratio,
             mlp_mult_factor=2,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
         )
 
     def forward(
@@ -389,6 +422,8 @@ class Flux2TransformerBlock(nn.Module):
         mlp_ratio: float = 3.0,
         eps: float = 1e-6,
         bias: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.mlp_hidden_dim = int(dim * mlp_ratio)
@@ -406,13 +441,29 @@ class Flux2TransformerBlock(nn.Module):
             added_proj_bias=bias,
             out_bias=bias,
             eps=eps,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
         )
 
         self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)
+        self.ff = Flux2FeedForward(
+            dim=dim,
+            dim_out=dim,
+            mult=mlp_ratio,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "ff"),
+        )
 
         self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.ff_context = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias)
+        self.ff_context = Flux2FeedForward(
+            dim=dim,
+            dim_out=dim,
+            mult=mlp_ratio,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "ff_context"),
+        )
 
     def forward(
         self,
@@ -576,6 +627,7 @@ class Flux2Transformer2DModel(nn.Module):
         rope_theta: int = 2000,
         eps: float = 1e-6,
         guidance_embeds: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
         self.guidance_embeds = guidance_embeds
@@ -624,8 +676,10 @@ class Flux2Transformer2DModel(nn.Module):
                     mlp_ratio=mlp_ratio,
                     eps=eps,
                     bias=False,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 
@@ -638,8 +692,10 @@ class Flux2Transformer2DModel(nn.Module):
                     mlp_ratio=mlp_ratio,
                     eps=eps,
                     bias=False,
+                    quant_config=quant_config,
+                    prefix=f"single_transformer_blocks.{i}",
                 )
-                for _ in range(num_single_layers)
+                for i in range(num_single_layers)
             ]
         )
 

vllm_omni/diffusion/models/flux2/pipeline_flux2.py

@@ -378,7 +378,7 @@ class Flux2Pipeline(nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarM
             self._execution_device
         )
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
-        self.transformer = Flux2Transformer2DModel(**transformer_kwargs)
+        self.transformer = Flux2Transformer2DModel(quant_config=od_config.quantization_config, **transformer_kwargs)
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
         self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -53,6 +53,10 @@ if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 
+def _join_prefix(prefix: str, name: str) -> str:
+    return f"{prefix}.{name}" if prefix else name
+
+
 class Flux2SwiGLU(nn.Module):
     """SwiGLU activation used by Flux2."""
 
@@ -74,6 +78,7 @@ class Flux2FeedForward(nn.Module):
         inner_dim: int | None = None,
         bias: bool = False,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         if inner_dim is None:
@@ -86,6 +91,7 @@ class Flux2FeedForward(nn.Module):
             bias=bias,
             return_bias=False,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "linear_in"),
         )
         self.act_fn = Flux2SwiGLU()
         self.linear_out = RowParallelLinear(
@@ -95,6 +101,7 @@ class Flux2FeedForward(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "linear_out"),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -119,6 +126,7 @@ class Flux2Attention(nn.Module):
         out_dim: int = None,
         elementwise_affine: bool = True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.parallel_config = parallel_config
@@ -136,6 +144,7 @@ class Flux2Attention(nn.Module):
             total_num_heads=self.heads,
             bias=bias,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "to_qkv"),
         )
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -152,6 +161,7 @@ class Flux2Attention(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
+                    prefix=_join_prefix(prefix, "to_out.0"),
                 ),
                 nn.Dropout(dropout),
             ]
@@ -166,6 +176,7 @@ class Flux2Attention(nn.Module):
                 total_num_heads=self.heads,
                 bias=added_proj_bias,
                 quant_config=quant_config,
+                prefix=_join_prefix(prefix, "add_kv_proj"),
             )
             self.add_query_num_heads = self.add_kv_proj.num_heads
             self.add_kv_num_heads = self.add_kv_proj.num_kv_heads
@@ -176,6 +187,7 @@ class Flux2Attention(nn.Module):
                 input_is_parallel=True,
                 return_bias=False,
                 quant_config=quant_config,
+                prefix=_join_prefix(prefix, "to_add_out"),
             )
 
         self.rope = RotaryEmbedding(is_neox_style=False)
@@ -330,6 +342,7 @@ class Flux2ParallelSelfAttention(nn.Module):
         mlp_ratio: float = 4.0,
         mlp_mult_factor: int = 2,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.parallel_config = parallel_config
@@ -350,6 +363,7 @@ class Flux2ParallelSelfAttention(nn.Module):
             bias=bias,
             gather_output=True,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "to_qkv_mlp_proj"),
         )
         self.mlp_act_fn = Flux2SwiGLU()
 
@@ -362,6 +376,7 @@ class Flux2ParallelSelfAttention(nn.Module):
             bias=out_bias,
             gather_output=True,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "to_out"),
         )
         self.rope = RotaryEmbedding(is_neox_style=False)
         self.attn = Attention(
@@ -463,6 +478,7 @@ class Flux2SingleTransformerBlock(nn.Module):
         eps: float = 1e-6,
         bias: bool = False,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -478,6 +494,7 @@ class Flux2SingleTransformerBlock(nn.Module):
             mlp_ratio=mlp_ratio,
             mlp_mult_factor=2,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
         )
 
     def forward(
@@ -534,6 +551,7 @@ class Flux2TransformerBlock(nn.Module):
         eps: float = 1e-6,
         bias: bool = False,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -551,10 +569,18 @@ class Flux2TransformerBlock(nn.Module):
             out_bias=bias,
             eps=eps,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
         )
 
         self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config)
+        self.ff = Flux2FeedForward(
+            dim=dim,
+            dim_out=dim,
+            mult=mlp_ratio,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "ff"),
+        )
 
         self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         self.ff_context = Flux2FeedForward(
@@ -563,6 +589,7 @@ class Flux2TransformerBlock(nn.Module):
             mult=mlp_ratio,
             bias=bias,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "ff_context"),
         )
 
     def forward(
@@ -835,8 +862,9 @@ class Flux2Transformer2DModel(nn.Module):
                     eps=eps,
                     bias=False,
                     quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
                 )
-                for _ in range(num_layers)
+                for i in range(num_layers)
             ]
         )
 
@@ -851,8 +879,9 @@ class Flux2Transformer2DModel(nn.Module):
                     eps=eps,
                     bias=False,
                     quant_config=quant_config,
+                    prefix=f"single_transformer_blocks.{i}",
                 )
-                for _ in range(num_single_layers)
+                for i in range(num_single_layers)
             ]
         )
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py

@@ -27,8 +27,16 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
 
 class HunyuanFusedMoEDefault(FusedMoE):
     def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
+        # Current vLLM FusedMoE handles output reduction internally.
+        kwargs.pop("reduce_results", None)
         super().__init__(prefix=prefix, **kwargs)
         self._prefix = prefix
+        self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True)
+
+    def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None:
+        if self.quant_method and getattr(self.quant_method, "moe_kernel", None) is None:
+            self.quant_method.process_weights_after_loading(self)
+        self._init_hook_handle.remove()
 
     def forward(self, hidden_states: Any, router_logits: Any) -> Any:
         _set_forward_context_num_tokens(hidden_states.shape[0])

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -2035,6 +2035,44 @@ class HunyuanImage3Model(nn.Module):
                     return True
             return False
 
+        def is_scalar_quant_scale(name: str, tensor: torch.Tensor) -> bool:
+            return tensor.numel() == 1 and name.endswith((".input_scale", ".weight_scale"))
+
+        def load_split_param(
+            name: str,
+            tensor: torch.Tensor,
+            den: int,
+            split_param: list[tuple[str | int, int]],
+            func: Callable[[torch.Tensor], torch.Tensor] | None,
+        ) -> None:
+            param = params_dict[name]
+            weight_loader = param.weight_loader
+            if is_scalar_quant_scale(name, tensor):
+                for shard_id, _ in split_param:
+                    weight_loader(param, tensor, shard_id)
+                return
+
+            assert tensor.shape[0] % den == 0
+            units = tensor.shape[0] // den
+            offset = 0
+            tensor = func(tensor) if func else tensor
+            for shard_id, num in split_param:
+                new_offset = offset + num * units
+                weight_loader(param, tensor[offset:new_offset], shard_id)
+                offset = new_offset
+
+        def get_loaded_weight_shard(
+            name: str,
+            tensor: torch.Tensor,
+            offset: int,
+            den: int,
+        ) -> torch.Tensor:
+            if is_scalar_quant_scale(name, tensor):
+                return tensor
+            assert tensor.shape[0] % den == 0
+            units = tensor.shape[0] // den
+            return tensor[offset * units : offset * units + units]
+
         for name, loaded_weight in weights:
             # print(f"Loading weight name: {name}, tp_rank: {tp_rank}", flush=True)
             if contains_unexpected_keyword(name, unexpected_keywords):
@@ -2112,19 +2150,7 @@ class HunyuanImage3Model(nn.Module):
                 if is_pp_missing_parameter(name, self):
                     continue
 
-                assert loaded_weight.shape[0] % den == 0
-                units = loaded_weight.shape[0] // den
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                offset = 0
-                for shard_id, num in split_param:
-                    new_offset = offset + num * units
-                    if func:
-                        weight_loader(param, func(loaded_weight)[offset:new_offset], shard_id)
-                    else:
-                        weight_loader(param, loaded_weight[offset:new_offset], shard_id)
-                    offset = new_offset
-
+                load_split_param(name, loaded_weight, den, split_param, func)
                 break
             else:
                 # Skip loading extra bias for GPTQ models.
@@ -2154,12 +2180,11 @@ class HunyuanImage3Model(nn.Module):
                         continue
                     param = params_dict[name_mapped]
                     weight_loader = cast(Callable[..., bool], param.weight_loader)
-                    assert loaded_weight.shape[0] % den == 0
-                    units = loaded_weight.shape[0] // den
+                    loaded_weight_shard = get_loaded_weight_shard(name, loaded_weight, offset, den)
 
                     success = weight_loader(
                         param,
-                        loaded_weight[offset * units : offset * units + units],
+                        loaded_weight_shard,
                         name_mapped,
                         shard_id=shard_id,
                         expert_id=expert_id,

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -209,6 +209,18 @@ def validate_zimage_tp_constraints(
     return ffn_hidden_dim, final_out_dims, supported_tp_candidates
 
 
+def _join_prefix(prefix: str, name: str) -> str:
+    return f"{prefix}.{name}" if prefix else name
+
+
+def _restore_linear_output_shape(output: torch.Tensor, input_: torch.Tensor) -> torch.Tensor:
+    if input_.dim() > 2 and output.dim() == 2:
+        token_count = math.prod(input_.shape[:-1])
+        if output.shape[0] == token_count:
+            return output.reshape(*input_.shape[:-1], output.shape[-1])
+    return output
+
+
 class TimestepEmbedder(nn.Module):
     def __init__(
         self, out_size, mid_size=None, frequency_embedding_size=256, quant_config: "QuantizationConfig | None" = None
@@ -256,6 +268,7 @@ class ZImageAttention(nn.Module):
         qk_norm: bool = True,
         eps: float = 1e-6,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.dim = dim
@@ -271,7 +284,7 @@ class ZImageAttention(nn.Module):
             total_num_kv_heads=num_kv_heads,
             bias=False,
             quant_config=quant_config,
-            prefix="to_qkv",
+            prefix=_join_prefix(prefix, "to_qkv"),
         )
 
         assert qk_norm is True
@@ -290,7 +303,7 @@ class ZImageAttention(nn.Module):
                     input_is_parallel=True,
                     return_bias=False,
                     quant_config=quant_config,
-                    prefix="to_out",
+                    prefix=_join_prefix(prefix, "to_out.0"),
                 )
             ]
         )
@@ -312,6 +325,7 @@ class ZImageAttention(nn.Module):
         sin: torch.Tensor,
     ):
         qkv, _ = self.to_qkv(hidden_states)
+        qkv = _restore_linear_output_shape(qkv, hidden_states)
         q_size = self.to_qkv.num_heads * self.head_dim
         kv_size = self.to_qkv.num_kv_heads * self.head_dim
         query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
@@ -344,7 +358,9 @@ class ZImageAttention(nn.Module):
         hidden_states = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.to(dtype)
 
+        to_out_input = hidden_states
         hidden_states = self.to_out[0](hidden_states)
+        hidden_states = _restore_linear_output_shape(hidden_states, to_out_input)
 
         return hidden_states
 
@@ -364,7 +380,7 @@ class FeedForward(nn.Module):
             bias=False,
             return_bias=False,
             quant_config=quant_config,
-            prefix=prefix,
+            prefix=_join_prefix(prefix, "w13"),
         )
         self.act = SiluAndMul()
         self.w2 = RowParallelLinear(
@@ -374,11 +390,15 @@ class FeedForward(nn.Module):
             input_is_parallel=True,
             return_bias=False,
             quant_config=quant_config,
-            prefix=prefix,
+            prefix=_join_prefix(prefix, "w2"),
         )
 
     def forward(self, x):
-        return self.w2(self.act(self.w13(x)))
+        hidden_states = self.w13(x)
+        hidden_states = _restore_linear_output_shape(hidden_states, x)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.w2(hidden_states)
+        return _restore_linear_output_shape(hidden_states, x)
 
 
 class ZImageTransformerBlock(nn.Module):
@@ -392,6 +412,7 @@ class ZImageTransformerBlock(nn.Module):
         qk_norm: bool,
         modulation=True,
         quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.dim = dim
@@ -403,10 +424,14 @@ class ZImageTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             eps=1e-5,
             quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attention"),
         )
 
         self.feed_forward = FeedForward(
-            dim=dim, hidden_dim=int(dim / 3 * 8), quant_config=quant_config, prefix="feed_forward"
+            dim=dim,
+            hidden_dim=int(dim / 3 * 8),
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "feed_forward"),
         )
         self.layer_id = layer_id
 
@@ -711,6 +736,7 @@ class ZImageTransformer2DModel(CachedTransformer):
                     qk_norm,
                     modulation=True,
                     quant_config=quant_config,
+                    prefix=f"noise_refiner.{layer_id}",
                 )
                 for layer_id in range(n_refiner_layers)
             ]
@@ -726,6 +752,7 @@ class ZImageTransformer2DModel(CachedTransformer):
                     qk_norm,
                     modulation=False,
                     quant_config=quant_config,
+                    prefix=f"context_refiner.{layer_id}",
                 )
                 for layer_id in range(n_refiner_layers)
             ]
@@ -757,6 +784,7 @@ class ZImageTransformer2DModel(CachedTransformer):
                     norm_eps,
                     qk_norm,
                     quant_config=quant_config,
+                    prefix=f"layers.{layer_id}",
                 )
                 for layer_id in range(n_layers)
             ]

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -11,9 +11,9 @@ to DiffusionModelRunner.
 import gc
 import multiprocessing as mp
 import os
-from collections.abc import Iterable
-from contextlib import AbstractContextManager, nullcontext
-from types import SimpleNamespace
+from collections.abc import Iterable, Iterator
+from contextlib import AbstractContextManager, contextmanager, nullcontext
+from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -22,7 +22,7 @@ from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
 from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
-from vllm.transformers_utils.config import get_config, get_hf_text_config
+from vllm.transformers_utils.config import get_hf_text_config
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.mem_utils import GiB_bytes
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -54,6 +54,78 @@ from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
 logger = init_logger(__name__)
 
 
+@dataclass
+class _DiffusionVllmModelConfig:
+    model: str
+    dtype: torch.dtype
+    quantization: str | None = None
+    quantization_config: Any | None = None
+    hf_config: Any | None = None
+    hf_text_config: Any | None = None
+    multimodal_config: Any | None = None
+    enforce_eager: bool = False
+    disable_cascade_attn: bool = False
+    enable_return_routed_experts: bool = False
+    is_moe: bool = False
+
+    def is_quantized(self) -> bool:
+        return self.quantization is not None
+
+    def is_model_moe(self) -> bool:
+        return self.is_moe
+
+    def is_nvfp4_quantized(self) -> bool:
+        return self.quantization == "modelopt_fp4"
+
+
+def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig:
+    quant_config = getattr(od_config, "quantization_config", None)
+    quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None
+    hf_config = getattr(od_config, "tf_model_config", None)
+    hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config
+    return _DiffusionVllmModelConfig(
+        model=od_config.model,
+        dtype=od_config.dtype,
+        quantization=quantization,
+        quantization_config=quant_config,
+        hf_config=hf_config,
+        hf_text_config=hf_text_config,
+        enforce_eager=getattr(od_config, "enforce_eager", False),
+        is_moe=bool(getattr(od_config, "is_moe", False)),
+    )
+
+
+@contextmanager
+def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[None]:
+    from vllm.model_executor.layers.quantization import modelopt as vllm_modelopt
+
+    linear_method_cls = getattr(quant_config, "LinearMethodCls", None)
+    if linear_method_cls in {
+        vllm_modelopt.ModelOptFp8LinearMethod,
+        vllm_modelopt.ModelOptFp8PcPtLinearMethod,
+    }:
+        from vllm.platforms import current_platform
+
+        if current_platform.is_cuda() and current_platform.has_device_capability(89):
+            from vllm.model_executor.kernels.linear import CutlassFP8ScaledMMLinearKernel
+
+            original_init_fp8_linear_kernel = vllm_modelopt.init_fp8_linear_kernel
+
+            def init_fp8_linear_kernel_with_cutlass(*args: Any, **kwargs: Any) -> Any:
+                kwargs.setdefault("force_kernel", CutlassFP8ScaledMMLinearKernel)
+                return original_init_fp8_linear_kernel(*args, **kwargs)
+
+            vllm_modelopt.init_fp8_linear_kernel = init_fp8_linear_kernel_with_cutlass
+            logger.info("Using CUTLASS FP8 linear kernels for this ModelOpt FP8 diffusion stage.")
+            try:
+                yield
+            finally:
+                vllm_modelopt.init_fp8_linear_kernel = original_init_fp8_linear_kernel
+            return
+
+    yield
+
+
 class DiffusionWorker:
     """
     A worker that manages GPU infrastructure and delegates to the model runner.
@@ -124,19 +196,7 @@ class DiffusionWorker:
         vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
         vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel
         vllm_config.profiler_config = self.od_config.profiler_config
-        try:
-            hf_config = get_config(self.od_config.model, trust_remote_code=self.od_config.trust_remote_code)
-        except ValueError:
-            hf_config = None
-            logger.info("Skipping hf_config loading for diffusion model %r", self.od_config.model_class_name)
-        hf_text_config = get_hf_text_config(hf_config) if hf_config is not None else None
-        vllm_config.model_config = SimpleNamespace(
-            hf_config=hf_config,
-            hf_text_config=hf_text_config,
-            enforce_eager=self.od_config.enforce_eager,
-            dtype=self.od_config.dtype,
-            enable_return_routed_experts=False,
-        )
+        vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config)  # type: ignore[assignment]
         vllm_config.quant_config = self.od_config.quantization_config
         # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping.
         # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion.
@@ -191,9 +251,15 @@ class DiffusionWorker:
         """Load the diffusion model using DiffusionModelRunner."""
         load_format = kwargs.get("load_format", load_format)
         custom_pipeline_name = kwargs.get("custom_pipeline_name", custom_pipeline_name)
+        cutlass_fp8_context = (
+            _force_cutlass_fp8_linear_kernel(self.od_config.quantization_config)
+            if getattr(self.od_config, "force_cutlass_fp8", False)
+            else nullcontext()
+        )
         with (
             set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config),
             set_current_vllm_config(self.vllm_config),
+            cutlass_fp8_context,
         ):
             self.model_runner.load_model(
                 memory_pool_context_fn=self._maybe_get_memory_pool_context,

vllm_omni/engine/arg_utils.py

@@ -140,6 +140,7 @@ class OmniEngineArgs(EngineArgs):
     async_chunk: bool = False
     omni_kv_config: dict | None = None
     quantization_config: Any | None = None
+    force_cutlass_fp8: bool | None = None
     worker_type: str | None = None
     task_type: str | None = None
     worker_cls: str = None

vllm_omni/engine/async_omni_engine.py

@@ -1467,6 +1467,7 @@ class AsyncOmniEngine:
             "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
             "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
             "quantization": kwargs.get("quantization", None),
+            "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)),
             "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
             **(

vllm_omni/entrypoints/cli/serve.py

@@ -338,6 +338,16 @@ class OmniServeCommand(CLISubcommand):
                 'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.'
             ),
         )
+        omni_config_group.add_argument(
+            "--force-cutlass-fp8",
+            action="store_true",
+            default=None,
+            help=(
+                "Diffusion-only runtime override for ModelOpt FP8 checkpoints: "
+                "force CUTLASS FP8 linear kernels on CUDA SM89+ devices. "
+                "Ignored for BF16, non-ModelOpt FP8, ROCm, and older CUDA GPUs."
+            ),
+        )
 
         # HSDP (Hybrid Sharded Data Parallel) parameters
         omni_config_group.add_argument(

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -132,7 +132,7 @@ class HunyuanModel(HunYuanModel):
             # Params for weights, fp8 weight scales, fp8 activation scales
             # (param_name, weight_name, expert_id, shard_id)
             fused_moe_expert_mapping = fused_moe_make_expert_params_mapping(
-                self,
+                model=self,
                 ckpt_gate_proj_name="gate_proj",
                 ckpt_down_proj_name="down_proj",
                 ckpt_up_proj_name="up_proj",

vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py

@@ -416,10 +416,8 @@ class BailingMoeV2SparseMoeBlock(nn.Module):
             dim=-1,
         )
 
-        # FusedMoE expects 2D hidden_states.
         hidden_states_2d = hidden_states.view(-1, h)
         final_hidden_states = self.experts(hidden_states_2d, packed_routing)
-
         final_hidden_states = final_hidden_states.view(bsz, seq_len, h)
 
         return final_hidden_states.squeeze(0) if input_is_2d else final_hidden_states

vllm_omni/quantization/factory.py

@@ -60,17 +60,83 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
     "int8": _build_int8,
     "inc": _build_inc,
     "auto-round": _build_inc,
+    "auto_round": _build_inc,
 }
 
 SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
 
 
+_MODEL_OPT_METHODS = {
+    "modelopt",
+}
+_MODEL_OPT_FP8_ALGOS = {
+    "FP8",
+    "FP8_PER_CHANNEL_PER_TOKEN",
+}
+
+
+def _normalize_method_name(method: Any) -> str:
+    return str(method).lower().replace("-", "_")
+
+
+def _detect_modelopt_method(config: Mapping[str, Any]) -> str | None:
+    quantization = config.get("quantization")
+    if isinstance(quantization, Mapping):
+        quant_algo = str(quantization.get("quant_algo", "")).upper()
+    else:
+        quant_algo = str(config.get("quant_algo", "")).upper()
+
+    method = config.get("method", config.get("quant_method"))
+    normalized_method = _normalize_method_name(method) if method is not None else None
+
+    producer = config.get("producer")
+    is_modelopt_config = normalized_method in _MODEL_OPT_METHODS or (
+        isinstance(producer, Mapping) and str(producer.get("name", "")).lower() == "modelopt"
+    )
+
+    if not is_modelopt_config:
+        return None
+
+    if quant_algo:
+        if quant_algo in _MODEL_OPT_FP8_ALGOS:
+            return "modelopt"
+        return None
+
+    if method is not None:
+        if normalized_method in _MODEL_OPT_METHODS:
+            return normalized_method
+
+    return None
+
+
+def _build_modelopt_from_config(method: str, config: Mapping[str, Any]) -> QuantizationConfig:
+    config_cls = get_quantization_config(method)
+    normalized_config = dict(config)
+    normalized_config.setdefault("quant_method", method)
+    return config_cls.from_config(normalized_config)
+
+
+def _pop_method_name(spec: dict[str, Any]) -> str | None:
+    method = spec.pop("method", None)
+    if method is None:
+        method = spec.pop("quant_method", None)
+    return method
+
+
+def _build_from_method_and_config(method: str, config: Mapping[str, Any]) -> QuantizationConfig:
+    normalized_config = {"quant_method": method, **config}
+    modelopt_method = _detect_modelopt_method(normalized_config)
+    if modelopt_method is not None:
+        return _build_modelopt_from_config(modelopt_method, normalized_config)
+    return _build_single(method, **config)
+
+
 def _build_single(method: str, **kwargs: Any) -> QuantizationConfig:
     """Build a single QuantizationConfig by method name.
 
     Resolution: _OVERRIDES first, then vLLM registry via from_config().
     """
-    method = method.lower()
+    method = _normalize_method_name(method)
 
     if method in _OVERRIDES:
         return _OVERRIDES[method](**kwargs)
@@ -92,16 +158,17 @@ def _build_single(method: str, **kwargs: Any) -> QuantizationConfig:
 def _is_per_component_dict(spec: dict[str, Any]) -> bool:
     """Check if a dict describes per-component quantization.
 
-    A per-component dict has no "method" key and all values are
+    A per-component dict has no "method" / "quant_method" key and all values are
     str, dict, or None. To avoid misdetecting a flat config with
     all-string values (e.g. {"activation_scheme": "static"}), we
-    require at least one value to be None or a dict with "method".
+    require at least one value to be None or a dict with "method" /
+    "quant_method".
     """
-    if "method" in spec:
+    if "method" in spec or "quant_method" in spec:
         return False
     if not all(isinstance(v, (dict, str, type(None))) for v in spec.values()):
         return False
-    return any(v is None or (isinstance(v, dict) and "method" in v) for v in spec.values())
+    return any(v is None or (isinstance(v, dict) and ("method" in v or "quant_method" in v)) for v in spec.values())
 
 
 def _build_component_config(spec: dict[str, Any]) -> ComponentQuantizationConfig:
@@ -116,10 +183,10 @@ def _build_component_config(spec: dict[str, Any]) -> ComponentQuantizationConfig
             config = _build_single(value)
         elif isinstance(value, dict):
             value = dict(value)  # avoid mutating caller's dict
-            method = value.pop("method", None)
+            method = _pop_method_name(value)
             if method is None:
-                raise ValueError(f"Component '{prefix}' config dict must have a 'method' key")
-            config = _build_single(method, **value)
+                raise ValueError(f"Component '{prefix}' config dict must have a 'method' or 'quant_method' key")
+            config = _build_from_method_and_config(method, value)
         else:
             raise TypeError(f"Component '{prefix}' config must be str, dict, or None, got {type(value).__name__}")
 
@@ -164,14 +231,19 @@ def build_quant_config(
         if _is_per_component_dict(spec):
             return _build_component_config(spec)
 
-        method = spec.pop("method", None)
+        modelopt_method = _detect_modelopt_method(spec)
+        if modelopt_method is not None:
+            logger.info("Building quantization config: %s", modelopt_method)
+            return _build_modelopt_from_config(modelopt_method, spec)
+
+        method = _pop_method_name(spec)
         if method is None:
             raise ValueError(
-                "Dict quantization config must have a 'method' key or "
+                "Dict quantization config must have a 'method' or 'quant_method' key or "
                 "be a per-component config with component prefixes as keys."
             )
         merged = {**spec, **kwargs}
         logger.info("Building quantization config: %s", method)
-        return _build_single(method, **merged)
+        return _build_from_method_and_config(method, merged)
 
     raise TypeError(f"quantization config must be str, dict, QuantizationConfig, or None, got {type(spec).__name__}")