[Core] Unified quantization framework (#1764)

benchmarks/diffusion/quantization_quality.py

@@ -0,0 +1,460 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Benchmark quantization quality loss for diffusion models (image & video).
+
+Generates outputs with BF16 (baseline) and a quantized config using the same
+seed, then computes LPIPS perceptual distance between them. Results are printed
+as a Markdown table ready to paste into a PR description.
+
+Requirements:
+    pip install lpips Pillow numpy
+
+Image example (text-to-image):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --task t2i \
+        --quantization fp8 \
+        --prompts \
+            "an aerial view of a coral reef with crystal clear turquoise water" \
+            "a campfire in a dark forest with sparks rising into a starry sky" \
+            "a gourmet dessert plate with chocolate mousse and gold leaf" \
+        --height 1024 --width 1024 \
+        --num-inference-steps 50 --seed 42
+
+Video example (text-to-video):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+        --task t2v \
+        --quantization fp8 \
+        --prompts \
+            "A serene lakeside sunrise with mist over the water" \
+            "A cat walking across a wooden bridge in autumn" \
+        --height 720 --width 1280 \
+        --num-frames 81 --num-inference-steps 40 --seed 42
+
+Multiple quantization methods:
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --task t2i \
+        --quantization fp8 int8 bitsandbytes \
+        --prompts "a cup of coffee on the table" \
+        --height 1024 --width 1024 \
+        --num-inference-steps 50 --seed 42
+
+Output directory structure (--output-dir, default: ./quant_bench_output):
+    quant_bench_output/
+        baseline/           # BF16 outputs
+        <method>/           # Quantized outputs per method
+        results.md          # Markdown table
+"""
+
+import argparse
+import gc
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def compute_lpips_images(
+    baseline_images: list,
+    quantized_images: list,
+    net: str = "alex",
+) -> list[float]:
+    """Compute LPIPS between paired lists of PIL images."""
+    import lpips
+    from torchvision import transforms
+
+    loss_fn = lpips.LPIPS(net=net).eval()
+    if torch.cuda.is_available():
+        loss_fn = loss_fn.cuda()
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize((256, 256)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+
+    scores = []
+    for img_bl, img_qt in zip(baseline_images, quantized_images):
+        t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
+        t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
+        if torch.cuda.is_available():
+            t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
+        with torch.no_grad():
+            score = loss_fn(t_bl, t_qt).item()
+        scores.append(score)
+    return scores
+
+
+def compute_lpips_video(
+    baseline_frames: np.ndarray,
+    quantized_frames: np.ndarray,
+    net: str = "alex",
+) -> float:
+    """Compute mean per-frame LPIPS for a video pair.
+
+    Args:
+        baseline_frames: (F, H, W, C) float array in [0, 1].
+        quantized_frames: same shape.
+
+    Returns:
+        Mean LPIPS across all frames.
+    """
+    import lpips
+
+    loss_fn = lpips.LPIPS(net=net).eval()
+    if torch.cuda.is_available():
+        loss_fn = loss_fn.cuda()
+
+    num_frames = min(len(baseline_frames), len(quantized_frames))
+    scores = []
+    for i in range(num_frames):
+        # Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
+        f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
+        f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
+        if torch.cuda.is_available():
+            f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
+        with torch.no_grad():
+            score = loss_fn(f_bl, f_qt).item()
+        scores.append(score)
+    return float(np.mean(scores))
+
+
+def _build_omni_kwargs(args, quantization=None):
+    """Build kwargs dict for Omni() constructor."""
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+    kwargs = {
+        "model": args.model,
+        "parallel_config": parallel_config,
+        "enforce_eager": args.enforce_eager,
+    }
+    if quantization:
+        kwargs["quantization_config"] = quantization
+    return kwargs
+
+
+def _generate_image(omni, args, prompt, seed):
+    """Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+    torch.cuda.reset_peak_memory_stats()
+    start = time.perf_counter()
+    outputs = omni.generate(
+        {"prompt": prompt},
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            num_inference_steps=args.num_inference_steps,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+    first = outputs[0]
+    req_out = first.request_output[0] if hasattr(first, "request_output") else first
+    img = req_out.images[0]
+    return img, elapsed, peak_mem
+
+
+def _generate_video(omni, args, prompt, seed):
+    """Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+    torch.cuda.reset_peak_memory_stats()
+    start = time.perf_counter()
+    outputs = omni.generate(
+        {"prompt": prompt, "negative_prompt": ""},
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+    first = outputs[0]
+    if hasattr(first, "request_output") and isinstance(first.request_output, list):
+        inner = first.request_output[0]
+        if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
+            frames = inner.images[0] if inner.images else None
+        else:
+            frames = inner
+    elif hasattr(first, "images") and first.images:
+        frames = first.images
+    else:
+        raise ValueError("Could not extract video frames from output.")
+
+    if isinstance(frames, torch.Tensor):
+        video = frames.detach().cpu()
+        if video.dim() == 5:
+            video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
+        elif video.dim() == 4 and video.shape[0] in (3, 4):
+            video = video.permute(1, 2, 3, 0)
+        if video.is_floating_point():
+            video = video.clamp(-1, 1) * 0.5 + 0.5
+        frames_array = video.float().numpy()
+    else:
+        frames_array = np.asarray(frames)
+        if frames_array.ndim == 5:
+            frames_array = frames_array[0]
+
+    return frames_array, elapsed, peak_mem
+
+
+def _unload_omni(omni):
+    """Delete Omni instance and free GPU memory."""
+    del omni
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+def run_benchmark(args):
+    from vllm_omni.entrypoints.omni import Omni
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    is_video = args.task == "t2v"
+    prompts = args.prompts
+    seed = args.seed
+
+    # Determine configs to benchmark
+    configs = []  # list of (label, quantization_method)
+    for method in args.quantization:
+        configs.append((method, method))
+
+    # --- Baseline run ---
+    print("\n" + "=" * 60)
+    print("Running BF16 baseline...")
+    print("=" * 60)
+    bl_kwargs = _build_omni_kwargs(args, quantization=None)
+    omni_bl = Omni(**bl_kwargs)
+
+    baseline_outputs = {}  # prompt -> (output, time, mem)
+    for prompt in prompts:
+        print(f"  Generating: {prompt[:60]}...")
+        if is_video:
+            out, t, mem = _generate_video(omni_bl, args, prompt, seed)
+        else:
+            out, t, mem = _generate_image(omni_bl, args, prompt, seed)
+        baseline_outputs[prompt] = (out, t, mem)
+
+    bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
+    bl_mem = baseline_outputs[prompts[0]][2]  # use first prompt's memory
+    _unload_omni(omni_bl)
+
+    # Save baseline outputs
+    bl_dir = output_dir / "baseline"
+    bl_dir.mkdir(parents=True, exist_ok=True)
+    for i, prompt in enumerate(prompts):
+        out = baseline_outputs[prompt][0]
+        if is_video:
+            try:
+                from diffusers.utils import export_to_video
+
+                frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
+                export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
+            except ImportError:
+                np.save(bl_dir / f"prompt_{i}.npy", out)
+        else:
+            out.save(bl_dir / f"prompt_{i}.png")
+
+    # --- Quantized runs ---
+    all_results = []  # list of dicts
+
+    for config_label, quant_method in configs:
+        print(f"\n{'=' * 60}")
+        print(f"Running: {config_label}...")
+        print("=" * 60)
+
+        qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
+        omni_qt = Omni(**qt_kwargs)
+
+        qt_outputs = {}
+        for prompt in prompts:
+            print(f"  Generating: {prompt[:60]}...")
+            if is_video:
+                out, t, mem = _generate_video(omni_qt, args, prompt, seed)
+            else:
+                out, t, mem = _generate_image(omni_qt, args, prompt, seed)
+            qt_outputs[prompt] = (out, t, mem)
+
+        qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
+        qt_mem = qt_outputs[prompts[0]][2]
+        _unload_omni(omni_qt)
+
+        # Save quantized outputs
+        qt_dir = output_dir / config_label.replace(" ", "_")
+        qt_dir.mkdir(parents=True, exist_ok=True)
+
+        # Compute LPIPS per prompt
+        per_prompt = []
+        for i, prompt in enumerate(prompts):
+            bl_out = baseline_outputs[prompt][0]
+            qt_out = qt_outputs[prompt][0]
+            if is_video:
+                lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
+                try:
+                    from diffusers.utils import export_to_video
+
+                    frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
+                    export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
+                except ImportError:
+                    np.save(qt_dir / f"prompt_{i}.npy", qt_out)
+            else:
+                lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
+                qt_out.save(qt_dir / f"prompt_{i}.png")
+            per_prompt.append({"prompt": prompt, "lpips": lpips_score})
+
+        mean_lpips = np.mean([p["lpips"] for p in per_prompt])
+        speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
+        mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
+
+        all_results.append(
+            {
+                "config": config_label,
+                "avg_time": qt_avg_time,
+                "speedup": speedup,
+                "memory_gib": qt_mem,
+                "mem_reduction_pct": mem_reduction,
+                "mean_lpips": mean_lpips,
+                "per_prompt": per_prompt,
+            }
+        )
+
+    # --- Print results ---
+    print("\n\n")
+    print("=" * 80)
+    print("RESULTS")
+    print("=" * 80)
+
+    # Summary table
+    lines = []
+    lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
+    lines.append(
+        f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
+        f"seed={args.seed}, LPIPS ({args.lpips_net})"
+    )
+    if is_video:
+        lines.append(f"Video: {args.num_frames} frames")
+    lines.append("")
+    lines.append("### Summary")
+    lines.append("")
+    lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
+    lines.append("|--------|----------|---------|--------------|---------------|------------|")
+    lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
+    for r in all_results:
+        lines.append(
+            f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
+            f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
+            f"| {r['mean_lpips']:.4f} |"
+        )
+    lines.append("")
+    lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
+    lines.append("")
+
+    # Per-prompt table
+    if len(prompts) > 1:
+        lines.append("### Per-Prompt LPIPS")
+        lines.append("")
+        header = "| Prompt |"
+        sep = "|--------|"
+        for r in all_results:
+            header += f" {r['config']} |"
+            sep += "--------|"
+        lines.append(header)
+        lines.append(sep)
+        for i, prompt in enumerate(prompts):
+            short = prompt[:50] + "..." if len(prompt) > 50 else prompt
+            row = f"| {short} |"
+            for r in all_results:
+                row += f" {r['per_prompt'][i]['lpips']:.4f} |"
+            lines.append(row)
+        lines.append("")
+
+    md = "\n".join(lines)
+    print(md)
+
+    # Save markdown
+    results_path = output_dir / "results.md"
+    results_path.write_text(md, encoding="utf-8")
+    print(f"\nResults saved to {results_path}")
+    print(f"Baseline outputs in {bl_dir}")
+    for r in all_results:
+        qt_dir = output_dir / r["config"].replace(" ", "_")
+        print(f"Quantized outputs in {qt_dir}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Benchmark quantization quality loss for diffusion models.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--model", required=True, help="Model name or local path.")
+    parser.add_argument(
+        "--task",
+        default="t2i",
+        choices=["t2i", "t2v"],
+        help="Task type: t2i (text-to-image) or t2v (text-to-video).",
+    )
+    parser.add_argument(
+        "--quantization",
+        nargs="+",
+        required=True,
+        help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
+    )
+    parser.add_argument(
+        "--prompts",
+        nargs="+",
+        default=["a cup of coffee on the table"],
+        help="One or more prompts to generate.",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
+    parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
+    parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
+    parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
+    parser.add_argument(
+        "--lpips-net",
+        type=str,
+        default="alex",
+        choices=["alex", "vgg", "squeeze"],
+        help="LPIPS backbone network.",
+    )
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--enforce-eager", action="store_true")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_benchmark(args)

docs/.nav.yml

@@ -79,6 +79,7 @@ nav:
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
+    - contributing/model/adding_quantization_model.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md

docs/contributing/model/adding_quantization_model.md

@@ -0,0 +1,478 @@
+# Support FP8 Quantization
+
+This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is FP8 Quantization?
+
+FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
+
+A diffusion pipeline typically has three quantizable components:
+
+| Component | Layer Types | Quantization Approach |
+|-----------|-----------|----------------------|
+| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
+| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
+| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
+
+### Architecture
+
+The quantization system has two mechanisms:
+
+**1. vLLM Quantized Linear Layers (for DiT)**
+
+vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
+
+Key classes:
+
+| Class | Location | Purpose |
+|-------|----------|---------|
+| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
+| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
+
+**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
+
+For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
+
+```
+At rest:     weight.data = fp8_tensor (half memory)
+Pre-hook:    weight.data = fp8_weight.to(bf16) * scale
+Forward:     Normal computation in BF16
+Post-hook:   weight.data = fp8_weight (back to FP8)
+```
+
+Key function:
+
+| Function | Location | Purpose |
+|----------|----------|---------|
+| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Add `quant_config` to the Transformer
+
+Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
+
+#### 1.1: Accept `quant_config` in Constructor
+
+```python
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+class YourTransformer2DModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        quant_config: QuantizationConfig | None = None,
+        # ... other params
+    ):
+        self.quant_config = quant_config
+        # Pass to sub-modules
+        self.blocks = nn.ModuleList([
+            YourTransformerBlock(quant_config=quant_config, ...)
+            for _ in range(num_layers)
+        ])
+```
+
+#### 1.2: Pass `quant_config` to vLLM Linear Layers
+
+Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
+
+```python
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+class YourAttentionBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
+        # QKV projection
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=hidden_size // num_heads,
+            total_num_heads=num_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_qkv",
+        )
+        # Output projection
+        self.to_out = RowParallelLinear(
+            input_size=hidden_size,
+            output_size=hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_out",
+        )
+
+class YourFeedForward(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
+        self.gate_up_proj = ColumnParallelLinear(
+            input_size=hidden_size,
+            output_size=intermediate_size * 2,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+```
+
+!!! important
+    Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
+
+#### 1.3: Identify Sensitive Layers
+
+Some layers degrade output quality when quantized. Common sensitive layers:
+
+| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
+|-------|-------------|---------------|----------------------|
+| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
+| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
+| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
+| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
+| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
+
+**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
+
+| Model | Method | All Layers | With `ignored_layers` | Recommendation |
+|-------|--------|-----------|----------------------|----------------|
+| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
+| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
+| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
+
+To identify sensitive layers for a new model:
+
+1. Run inference with all layers quantized
+2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
+3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
+4. Document the recommended `ignored_layers` in the model's supported models table
+
+### Step 2: Wire Up Quantization in the Pipeline
+
+#### 2.1: Create `quant_config` and Pass to Transformer
+
+In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+class YourPipeline(nn.Module):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        # ... load scheduler, text_encoder, vae ...
+
+        # Build quant config from OmniDiffusionConfig
+        quant_config = build_quant_config(od_config.quantization_config)
+
+        # Pass to transformer
+        self.transformer = YourTransformer2DModel(
+            od_config=od_config,
+            quant_config=quant_config,
+            # ... other kwargs
+        )
+```
+
+#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
+
+For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
+
+!!! note
+    The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
+
+```python
+from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
+
+class YourPipeline(nn.Module):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+
+        self.text_encoder = SomeTextEncoder.from_pretrained(...)
+        self.vae = SomeVAE.from_pretrained(...).to(self.device)
+
+        # Apply FP8 weight storage to text encoder and VAE
+        if (
+            od_config.quantization_config is not None
+            and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
+        ):
+            apply_fp8_weight_storage(self.vae)
+            apply_fp8_weight_storage(self.text_encoder)
+
+        # ... rest of init
+```
+
+`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
+
+#### 2.3: Fix `load_weights` to Include Pre-loaded Components
+
+Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
+
+```python
+def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    loader = AutoWeightsLoader(self)
+    loaded_weights = loader.load_weights(weights)
+    # Mark pre-loaded component weights as loaded
+    loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
+    loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
+    return loaded_weights
+```
+
+### Step 3: Add to Supported Models Documentation
+
+Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
+
+```markdown
+| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
+```
+
+---
+
+## Testing
+
+### Quick Verification
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Test with FP8 enabled
+omni = Omni(model="your-model", quantization="fp8")
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+### Quality Gate Test (LPIPS)
+
+We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
+It generates outputs with both BF16 and your quantized method using the same seed, computes
+[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
+exceeds a threshold**. This is the recommended way to validate that a quantization method does not
+unacceptably degrade output quality.
+
+**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
+
+```python
+QualityTestConfig(
+    id="int8_z_image",              # pytest ID
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization="int8",            # your quantization method
+    task="t2i",                     # "t2i" or "t2v"
+    prompt="a cup of coffee on a wooden table, morning light",
+    max_lpips=0.15,                 # fail threshold (see guide below)
+    num_inference_steps=20,
+),
+```
+
+**Step 2: Run the test:**
+
+```bash
+# Install dependency
+pip install lpips
+
+# Run all quality tests
+pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
+
+# Run only your method
+pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
+
+# Run a specific model
+pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
+```
+
+**Step 3: Read the output:**
+
+```
+Quantization Quality: int8_z_image
+  Model:         Tongyi-MAI/Z-Image-Turbo
+  Method:        int8
+  LPIPS:         0.0290  (threshold: 0.15)
+  BF16 memory:   19.15 GiB
+  Quant memory:  14.80 GiB  (23% reduction)
+  Result:        PASS
+```
+
+**LPIPS threshold guide:**
+
+| LPIPS Range | Meaning | Recommendation |
+|-------------|---------|----------------|
+| < 0.01 | Imperceptible | Ideal — no visible difference |
+| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
+| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
+| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
+| > 0.30 | Significant degradation | Quantization method may not suit this model |
+
+!!! tip
+    If your test fails, don't raise the threshold — instead, identify sensitive layers and add
+    them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
+
+### Quality Benchmark (Detailed)
+
+For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
+standalone benchmark script:
+
+```bash
+python benchmarks/diffusion/quantization_quality.py \
+    --model Tongyi-MAI/Z-Image-Turbo \
+    --quantization fp8 \
+    --prompts \
+        "a cup of coffee on the table" \
+        "an aerial view of a coral reef" \
+    --height 1024 --width 1024 \
+    --num-inference-steps 50 --seed 42
+```
+
+This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
+summary table ready to paste into your PR.
+
+### Identifying Sensitive Layers with `ignored_layers`
+
+Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
+in BF16 while quantizing everything else. This is critical for achieving good quality with
+aggressive quantization methods (Int8, NVFP4).
+
+**Why some layers are sensitive:**
+
+| Layer Type | Why Sensitive | Typical Impact |
+|------------|--------------|----------------|
+| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
+| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
+| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
+| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
+| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
+
+**How to find sensitive layers for your model:**
+
+1. Run the quality gate test with all layers quantized:
+    ```bash
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
+    ```
+
+2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
+    ```python
+    QualityTestConfig(
+        id="fp8_your_model_skip_mlp",
+        model="your-org/Your-Model",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table",
+        max_lpips=0.05,
+    ),
+    ```
+
+    And in the test, modify the Omni call:
+    ```python
+    omni_qt = Omni(
+        model=config.model,
+        quantization_config={
+            "method": config.quantization,
+            "ignored_layers": ["img_mlp"],
+        },
+    )
+    ```
+
+3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
+
+4. Document the recommended `ignored_layers` in the supported models table.
+
+**Real-world examples:**
+
+| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
+|-------|--------|-------------------|----------------------|---------------------|
+| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
+| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
+| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
+
+### What to Verify (Checklist)
+
+1. **Model loads without errors** — Check logs for quantization initialization messages
+2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
+3. **Memory reduction** — Expect ~30-50% savings on quantized components
+4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
+5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
+
+---
+
+## Troubleshooting
+
+### Issue: Output Quality Degradation
+
+**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
+
+**Causes & Solutions:**
+
+- **Sensitive layers being quantized:**
+
+    **Solution:** Identify and skip sensitive layers:
+    ```python
+    omni = Omni(
+        model="your-model",
+        quantization_config={
+            "method": "fp8",
+            "ignored_layers": ["img_mlp"],
+        },
+    )
+    ```
+
+- **All layers sensitive (rare):**
+
+    **Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
+
+### Issue: Missing Weights Warning
+
+**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
+
+**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
+
+### Issue: CUDA Out of Memory
+
+**Symptoms:** OOM errors during model loading despite FP8 being enabled.
+
+**Causes & Solutions:**
+
+- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
+
+    FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
+
+---
+
+## Reference Implementations
+
+| Model | Pipeline File | Transformer File | Notes |
+|-------|-------------|-----------------|-------|
+| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
+| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
+
+All files are under `vllm_omni/diffusion/`.
+
+---
+
+## Summary
+
+Adding FP8 quantization support to a new model:
+
+1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
+2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
+3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
+4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
+5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
+6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior

docs/mkdocs/hooks/generate_api_readme.py

@@ -139,6 +139,8 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
             # Skip excluded modules (avoid importing vllm during docs build)
             excluded_prefixes = [
                 "vllm_omni.diffusion.models.qwen_image",
+                "vllm_omni.diffusion.quantization",
+                "vllm_omni.quantization",
                 "vllm_omni.entrypoints.async_diffusion",
                 "vllm_omni.entrypoints.openai",
                 "vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",

docs/user_guide/diffusion/quantization/fp8.md

@@ -63,6 +63,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
 | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
 | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
+| HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
 
 ## Combining with Other Features

docs/user_guide/diffusion/quantization/gguf.md

@@ -12,7 +12,7 @@
 
 ## Architecture Overview
 1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
-2. Diffusion quantization wrapper (`DiffusionGgufConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
+2. Diffusion quantization wrapper (`DiffusionGGUFConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
 3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
 4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
 5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.

docs/user_guide/diffusion/quantization/overview.md

@@ -1,22 +1,86 @@
-# Quantization for Diffusion Transformers
+# Quantization
 
-vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and accelerate inference.
+vLLM-Omni provides a unified quantization framework that supports both diffusion models (DiT, text encoders, VAEs) and multi-stage omni models (e.g., Qwen3-Omni thinker). Quantization reduces memory usage and accelerates inference.
 
 ## Supported Methods
 
-| Method | Guide |
-|--------|-------|
-| FP8 | [FP8](fp8.md) |
-| Int8 | [Int8](int8.md) |
-| GGUF | [GGUF](gguf.md) |
+### Diffusion Models
+
+| Method | Guide | Description | Tested Models | Min GPU |
+|--------|-------|-------------|---------------|---------|
+| FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
+| Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
+| GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
+
+### Multi-stage Omni Models (Pre-quantized Checkpoints)
+
+| Method | Description | Tested Models | Min GPU |
+|--------|-------------|---------------|---------|
+| ModelOpt FP8 | Pre-quantized FP8 via NVIDIA ModelOpt | Qwen3-Omni (thinker) | SM 89 (Ada/Hopper) |
+| ModelOpt NVFP4 | Pre-quantized NVFP4 via NVIDIA ModelOpt | Qwen3-Omni (experimental, quality issues) | SM 100 (Blackwell) |
+
+!!! note
+    AWQ, GPTQ, and BitsAndBytes are supported by vLLM's upstream quantization registry but have **not been tested** with vLLM-Omni pipelines. They may work via `build_quant_config()` but are not validated.
+
+### Pre-quantized LLM Checkpoints (Multi-stage Models)
+
+For multi-stage models like Qwen3-Omni, the unified quantization framework auto-detects
+pre-quantized checkpoints via `quantization_config` in the HF config. Quantization is
+automatically scoped to the thinker's `language_model` — audio encoder, vision encoder,
+talker, and code2wav remain in BF16.
+
+| Format | `quant_algo` | Hardware | Example | Status |
+|--------|-------------|----------|---------|--------|
+| ModelOpt FP8 | `FP8` | Ada/Hopper (SM 89+) | `asdazd/Qwen3-Omni-30B-A3B-Instruct_modelopt_FP8` | Tested — 47% memory reduction, comparable throughput |
+| ModelOpt NVFP4 | `NVFP4` | Blackwell (SM 100+) | `shunyang90/Qwen3-Omni-30B-A3B-Instruct-NVFP4` | Experimental — loads and runs but output quality is unacceptable |
+
+**Tested FP8 results (Qwen3-Omni, 1×H200):**
+
+| Config | Model Memory (GiB) | Mem Reduction | Decode (tok/s) |
+|--------|-------------------|---------------|----------------|
+| BF16 baseline | 59.26 | — | 41.6 |
+| FP8 (ModelOpt) | 31.41 | 47% | 39.9 |
+
+FP8 enables the full 3-stage pipeline (thinker + talker + code2wav) on a single 64GB GPU,
+which is impossible with BF16 (thinker alone requires 59.26 GiB).
+
+## Quantization Scope
+
+### Diffusion Models
+
+When `--quantization fp8` is enabled for diffusion models:
+
+| Component | What Gets Quantized | Mechanism |
+|-----------|-------------------|-----------|
+| **DiT (transformer)** | `nn.Linear` layers | vLLM W8A8 FP8 compute (Ada/Hopper) or weight-only (older GPUs) |
+| **Text encoder** | `nn.Linear` layers | FP8 weight storage, BF16 compute |
+| **VAE** | `nn.Conv2d`, `nn.Conv3d` layers | FP8 weight storage, BF16 compute |
+
+### Multi-stage Omni Models
+
+For pre-quantized omni model checkpoints, quantization is scoped per component:
+
+| Component | Quantized? | Notes |
+|-----------|-----------|-------|
+| **Thinker (language_model)** | Yes | FP8 or NVFP4 via ModelOpt |
+| **Audio encoder** | No | Stays BF16 |
+| **Vision encoder** | No | Stays BF16 |
+| **Talker** | No | Stays BF16 |
+| **Code2Wav** | No | Stays BF16 |
+
+!!! note
+    Not all models support all components. See the [FP8 supported models table](fp8.md#supported-models) for per-model details.
 
 ## Device Compatibility for FP8
 
-| GPU Generation | Example GPUs | FP8 Mode |
-|---------------|-------------------|----------|
-| Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
+| Device | Example Hardware | FP8 | NVFP4 |
+|--------|-----------------|-----|-------|
+| Blackwell GPU (SM 100+) | RTX 5090, B100, B200 | Yes | Yes (native FP4 HW) |
+| Ada/Hopper GPU (SM 89+) | RTX 4090, H100, H200 | Yes (W8A8 native) | No |
+| Turing/Ampere GPU (SM 75-86) | RTX 3090, A100 | Yes (weight-only Marlin) | No |
+| Ascend NPU | Atlas 800T A2 (910B) | Not yet supported | No |
 
-Kernel selection is automatic.
+Kernel selection is automatic on CUDA GPUs.
 
 ## Device Compatibility for Int8
 
@@ -24,3 +88,62 @@ Kernel selection is automatic.
 |-------------|---------------|-------------------|----------|
 | NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
 | Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |
+
+## Python API
+
+The `build_quant_config()` factory accepts multiple input formats:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+# String
+config = build_quant_config("fp8")
+
+# Dict with parameters
+config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
+
+# Per-component dict
+config = build_quant_config({
+    "transformer": {"method": "fp8"},
+    "vae": None,
+})
+
+# None / "none"
+config = build_quant_config(None)  # returns None
+```
+
+## Architecture
+
+```
+build_quant_config(spec)
+    ├── str  → _build_single(method) → vLLM registry / _OVERRIDES
+    ├── dict with "method" → _build_single(method, **kwargs)
+    ├── per-component dict → ComponentQuantizationConfig
+    │       routes get_quant_method() by longest-prefix match
+    ├── QuantizationConfig → passthrough
+    └── None / "none" → None
+```
+
+## Migration Guide
+
+### Before (v0.14.x)
+
+```python
+# Old diffusion-specific API (removed)
+from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+config = get_diffusion_quant_config("fp8", activation_scheme="static")
+```
+
+### After (v0.16.0+)
+
+```python
+# Unified API — delegates to vLLM's registry
+from vllm_omni.quantization import build_quant_config
+config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
+
+# Per-component (new)
+config = build_quant_config({
+    "transformer": {"method": "fp8"},
+    "vae": None,
+})
+```

examples/offline_inference/bagel/end2end.py

@@ -70,6 +70,12 @@ def parse_args():
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        help="Quantization method (e.g. 'fp8').",
+    )
 
     args = parser.parse_args()
     return args
@@ -119,6 +125,8 @@ def main():
             "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         }
     )
+    if args.quantization:
+        omni_kwargs["quantization_config"] = args.quantization
 
     omni = Omni(model=model_name, **omni_kwargs)
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -263,7 +263,7 @@ query_map = {
 
 
 def main(args):
-    model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    model_name = args.model
     print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
 
     # Get paths from args
@@ -299,6 +299,7 @@ def main(args):
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
+        init_timeout=args.init_timeout,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
@@ -333,11 +334,14 @@ def main(args):
         repetition_penalty=1.1,
     )
 
-    sampling_params_list = [
+    all_sampling_params = [
         thinker_sampling_params,
         talker_sampling_params,  # code predictor is integrated into talker for Qwen3 Omni
         code2wav_sampling_params,
     ]
+    # Match sampling params to the number of configured stages
+    num_stages = omni.num_stages
+    sampling_params_list = all_sampling_params[:num_stages]
 
     if args.txt_prompts is None:
         prompts = [query_result.inputs for _ in range(args.num_prompts)]
@@ -415,6 +419,12 @@ def main(args):
 
 def parse_args():
     parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        help="Model name or path.",
+    )
     parser.add_argument(
         "--query-type",
         "-q",

mkdocs.yml

@@ -98,6 +98,7 @@ plugins:
         - "re:vllm_omni\\._.*"  # Internal modules
         - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building
         - "vllm_omni.diffusion.quantization"  # avoid importing vllm in mkdocs building
+        - "vllm_omni.quantization"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.async_diffusion"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.openai"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.openai.protocol"  # avoid importing vllm in mkdocs building

tests/diffusion/quantization/test_fp8_config.py

@@ -1,137 +1,366 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for FP8 quantization config."""
+"""Tests for the unified quantization framework."""
 
 import pytest
+from torch import nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 
-def test_fp8_config_creation():
-    """Test that FP8 config can be created."""
-    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+def test_build_quant_config_fp8():
+    from vllm_omni.quantization import build_quant_config
 
-    config = get_diffusion_quant_config("fp8")
+    config = build_quant_config("fp8")
     assert config is not None
     assert config.get_name() == "fp8"
+    assert config.activation_scheme == "dynamic"
 
 
-def test_vllm_config_extraction():
-    """Test that vLLM config can be extracted from diffusion config."""
-    from vllm_omni.diffusion.quantization import (
-        get_diffusion_quant_config,
-        get_vllm_quant_config_for_layers,
-    )
+def test_build_quant_config_none():
+    from vllm_omni.quantization import build_quant_config
 
-    diff_config = get_diffusion_quant_config("fp8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
-    assert vllm_config is not None
-    assert vllm_config.activation_scheme == "dynamic"
+    assert build_quant_config(None) is None
 
 
-def test_none_quantization():
-    """Test that None quantization returns None config."""
-    from vllm_omni.diffusion.quantization import (
-        get_diffusion_quant_config,
-        get_vllm_quant_config_for_layers,
-    )
+def test_build_quant_config_none_string():
+    from vllm_omni.quantization import build_quant_config
 
-    config = get_diffusion_quant_config(None)
-    assert config is None
-    vllm_config = get_vllm_quant_config_for_layers(config)
-    assert vllm_config is None
+    assert build_quant_config("none") is None
 
 
-def test_invalid_quantization():
-    """Test that invalid quantization method raises error."""
-    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+def test_build_quant_config_invalid():
+    from vllm_omni.quantization import build_quant_config
 
     with pytest.raises(ValueError, match="Unknown quantization method"):
-        get_diffusion_quant_config("invalid_method")
+        build_quant_config("invalid_method")
 
 
-def test_fp8_config_with_custom_params():
-    """Test FP8 config with custom parameters."""
-    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+def test_build_quant_config_dict():
+    from vllm_omni.quantization import build_quant_config
 
-    config = get_diffusion_quant_config(
-        "fp8",
-        activation_scheme="static",
-        ignored_layers=["proj_out"],
-    )
+    config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
     assert config is not None
+    assert config.get_name() == "fp8"
     assert config.activation_scheme == "static"
-    assert "proj_out" in config.ignored_layers
 
 
-def test_supported_methods():
-    """Test that supported methods list is correct."""
-    from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
+def test_build_quant_config_dict_not_mutated():
+    from vllm_omni.quantization import build_quant_config
 
-    assert "fp8" in SUPPORTED_QUANTIZATION_METHODS
+    original = {"method": "fp8", "activation_scheme": "static"}
+    copy = original.copy()
+    build_quant_config(original)
+    assert original == copy
 
 
-def test_quantization_integration():
-    """Test end-to-end quantization flow through OmniDiffusionConfig."""
+def test_build_quant_config_per_component():
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config(
+        {
+            "transformer": {"method": "fp8"},
+            "vae": None,
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+    assert config.component_configs["transformer"].get_name() == "fp8"
+    assert config.component_configs["vae"] is None
+
+
+def test_build_quant_config_per_component_string():
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config({"transformer": "fp8", "vae": None})
+    assert isinstance(config, ComponentQuantizationConfig)
+    assert config.component_configs["transformer"].get_name() == "fp8"
+
+
+def test_build_quant_config_per_component_inner_dict_not_mutated():
+    """Inner component dicts should not be mutated by build_quant_config."""
+    from vllm_omni.quantization import build_quant_config
+
+    inner = {"method": "fp8", "activation_scheme": "static"}
+    original = inner.copy()
+    build_quant_config({"transformer": inner, "vae": None})
+    assert inner == original
+
+
+def test_flat_dict_not_misdetected_as_per_component():
+    """A flat config like {"activation_scheme": "static"} must NOT be treated
+    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"):
+        build_quant_config({"activation_scheme": "static"})
+
+
+def test_build_quant_config_passthrough():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import build_quant_config
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    assert build_quant_config(fp8) is fp8
+
+
+def test_component_config_routing():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    config = ComponentQuantizationConfig(
+        component_configs={"transformer": fp8, "vae": None},
+    )
+
+    assert config.get_name() == "component"
+    assert config.resolve("transformer.blocks.0.attn") is fp8
+    assert config.resolve("vae.encoder.conv_in") is None
+    assert config.resolve("unknown.layer") is None
+
+
+def test_component_config_with_default():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.quantization import ComponentQuantizationConfig
+
+    fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
+    config = ComponentQuantizationConfig(
+        component_configs={"vae": None},
+        default_config=fp8,
+    )
+
+    assert config.resolve("transformer.blocks.0") is fp8
+    assert config.resolve("vae.encoder") is None
+
+
+def test_gguf_config():
+    from vllm_omni.quantization import build_quant_config
+    from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
+
+    config = build_quant_config(
+        {
+            "method": "gguf",
+            "gguf_model": "path/to/model.gguf",
+        }
+    )
+    assert isinstance(config, DiffusionGGUFConfig)
+    assert config.gguf_model == "path/to/model.gguf"
+    assert config.get_name() == "gguf"
+
+
+def test_integration_string():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    # Test with quantization string only
-    config = OmniDiffusionConfig(model="test", quantization="fp8")
+    config = OmniDiffusionConfig(model="test", quantization_config="fp8")
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "fp8"
 
-    # Test with quantization_config dict
-    config2 = OmniDiffusionConfig(
+
+def test_integration_dict():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig(
         model="test",
         quantization_config={"method": "fp8", "activation_scheme": "static"},
     )
-    assert config2.quantization_config is not None
-    assert config2.quantization_config.get_name() == "fp8"
-    assert config2.quantization_config.activation_scheme == "static"
-
-    # Test that vLLM config can be extracted
-    vllm_config = config.quantization_config.get_vllm_quant_config()
-    assert vllm_config is not None
-
-
-def test_quantization_dict_not_mutated():
-    """Test that passing a dict to quantization_config doesn't mutate it."""
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-    original_dict = {"method": "fp8", "activation_scheme": "static"}
-    dict_copy = original_dict.copy()
-
-    OmniDiffusionConfig(model="test", quantization_config=original_dict)
-
-    # Original dict should be unchanged
-    assert original_dict == dict_copy
-
-
-def test_quantization_conflicting_methods_warning(caplog):
-    """Test warning when quantization and quantization_config['method'] conflict."""
-    import logging
-
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-    with caplog.at_level(logging.WARNING):
-        config = OmniDiffusionConfig(
-            model="test",
-            quantization="fp8",  # This should be overridden
-            quantization_config={"method": "fp8", "activation_scheme": "static"},
-        )
-    # No warning when methods match
     assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "fp8"
+    assert config.quantization_config.activation_scheme == "static"
 
 
-def test_fp8_delegates_to_vllm_config():
-    """Test that DiffusionFp8Config delegates to vLLM's Fp8Config."""
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+def test_integration_no_quant():
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    from vllm_omni.diffusion.quantization import DiffusionFp8Config
+    config = OmniDiffusionConfig(model="test")
+    assert config.quantization_config is None
 
-    # Test that quant_config_cls is set correctly
-    assert DiffusionFp8Config.quant_config_cls is Fp8Config
 
-    # Test that get_name() delegates to vLLM
-    assert DiffusionFp8Config.get_name() == Fp8Config.get_name()
+def test_integration_per_component():
+    """OmniDiffusionConfig with per-component quantization dict."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.quantization import ComponentQuantizationConfig
 
-    # Test that get_min_capability() delegates to vLLM
-    assert DiffusionFp8Config.get_min_capability() == Fp8Config.get_min_capability()
+    config = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "transformer": {"method": "fp8"},
+            "vae": None,
+        },
+    )
+    assert isinstance(config.quantization_config, ComponentQuantizationConfig)
+    assert config.quantization_config.component_configs["transformer"].get_name() == "fp8"
+    assert config.quantization_config.component_configs["vae"] is None
+
+
+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"]:
+        assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
+
+
+def test_supported_methods_count():
+    from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert len(SUPPORTED_QUANTIZATION_METHODS) >= 20
+
+
+def test_per_component_routing_with_resolve():
+    """Verify resolve() routes correctly by prefix."""
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config(
+        {
+            "transformer": {"method": "fp8"},
+            "vae": None,
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+
+    assert config.resolve("transformer.blocks.0.attn.to_q") is not None
+    assert config.resolve("transformer.blocks.0.attn.to_q").get_name() == "fp8"
+    assert config.resolve("vae.encoder.conv_in") is None
+    assert config.resolve("unknown.layer.0.weight") is None
+
+
+def test_per_component_routing_with_default():
+    """Verify default config applies to unmatched prefixes."""
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    config = build_quant_config(
+        {
+            "vae": None,
+            "default": "fp8",
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+
+    assert config.resolve("vae.decoder.conv") is None
+    resolved = config.resolve("transformer.blocks.0.attn")
+    assert resolved is not None
+    assert resolved.get_name() == "fp8"
+
+
+@pytest.mark.parametrize(
+    "quant_algo",
+    ["FP8", "NVFP4"],
+    ids=["modelopt_fp8", "modelopt_nvfp4"],
+)
+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."""
+    from types import SimpleNamespace
+
+    from vllm_omni.config.model import OmniModelArchConfigConvertor
+
+    text_config = SimpleNamespace(
+        quantization_config={
+            "quant_method": "modelopt",
+            "quant_algo": quant_algo,
+            "ignore": ["lm_head", "model.layers.0.mlp.gate"],
+        },
+        model_type="qwen3_moe",
+    )
+    thinker_config = SimpleNamespace(text_config=text_config)
+    hf_config = SimpleNamespace(
+        thinker_config=thinker_config,
+        talker_config=SimpleNamespace(text_config=SimpleNamespace()),
+        model_type="qwen3_omni_moe",
+    )
+
+    convertor = OmniModelArchConfigConvertor(hf_config, text_config, stage_config_name="thinker_config")
+    quant_cfg = convertor.get_quantization_config()
+
+    assert quant_cfg is not None
+    assert quant_cfg["quant_method"] == "modelopt"
+    assert "lm_head" in quant_cfg["ignore"]
+
+
+def test_omni_convertor_talker_returns_none():
+    """Talker stage should get no quantization config when its text_config
+    has no quantization_config (talker weights are BF16)."""
+    from types import SimpleNamespace
+
+    from vllm_omni.config.model import OmniModelArchConfigConvertor
+
+    talker_text_config = SimpleNamespace(model_type="qwen3_omni_moe_talker")
+    talker_config = SimpleNamespace(text_config=talker_text_config)
+    hf_config = SimpleNamespace(
+        talker_config=talker_config,
+        model_type="qwen3_omni_moe",
+    )
+
+    convertor = OmniModelArchConfigConvertor(hf_config, talker_text_config, stage_config_name="talker_config")
+    quant_cfg = convertor.get_quantization_config()
+
+    assert quant_cfg is None
+
+
+def test_omni_convertor_no_stage_name_falls_back():
+    """Without stage_config_name, should fall back to base behavior."""
+    from types import SimpleNamespace
+
+    from vllm_omni.config.model import OmniModelArchConfigConvertor
+
+    hf_config = SimpleNamespace(model_type="qwen3_omni_moe")
+    text_config = SimpleNamespace()
+
+    convertor = OmniModelArchConfigConvertor(hf_config, text_config)
+    quant_cfg = convertor.get_quantization_config()
+
+    assert quant_cfg is None
+
+
+def test_multi_component_model_routing():
+    """Integration test: walk a multi-component model and verify per-component
+    quantization routes resolve() correctly for every linear layer."""
+    from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config
+
+    # Build a mock multi-stage model mimicking Bagel/Qwen3-Omni layout
+    class MockTransformerBlock(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.attn_q = nn.Linear(64, 64)
+            self.attn_k = nn.Linear(64, 64)
+            self.mlp = nn.Linear(64, 256)
+
+    class MockVAEBlock(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = nn.Linear(64, 64)
+
+    class MockMultiStageModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.transformer = nn.ModuleDict(
+                {
+                    "block_0": MockTransformerBlock(),
+                    "block_1": MockTransformerBlock(),
+                }
+            )
+            self.vae = nn.ModuleDict(
+                {
+                    "encoder": MockVAEBlock(),
+                    "decoder": MockVAEBlock(),
+                }
+            )
+
+    model = MockMultiStageModel()
+    config = build_quant_config(
+        {
+            "transformer": {"method": "fp8", "activation_scheme": "dynamic"},
+            "vae": None,
+        }
+    )
+    assert isinstance(config, ComponentQuantizationConfig)
+
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear):
+            resolved = config.resolve(name)
+            if name.startswith("transformer"):
+                assert resolved is not None, f"{name} should be quantized"
+                assert resolved.get_name() == "fp8"
+            elif name.startswith("vae"):
+                assert resolved is None, f"{name} should NOT be quantized"

tests/diffusion/quantization/test_gguf_config.py

@@ -9,16 +9,12 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm.model_executor.layers.linear import LinearBase
-from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES
+from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES, UnquantizedLinearMethod
 
-from vllm_omni.diffusion.quantization import (
-    DiffusionGgufConfig,
-    get_diffusion_quant_config,
-)
-from vllm_omni.diffusion.quantization.gguf import (
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.gguf_config import (
+    DiffusionGGUFConfig,
     DiffusionGGUFLinearMethod,
-    UnquantizedLinearMethod,
-    _GGUFConfig,
     dequant_gemm_gguf,
 )
 
@@ -26,7 +22,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def test_gguf_config_creation_and_delegation():
-    config = DiffusionGgufConfig(
+    config = DiffusionGGUFConfig(
         gguf_model="weights.gguf",
         unquantized_modules=["proj_out"],
     )
@@ -34,37 +30,36 @@ def test_gguf_config_creation_and_delegation():
     assert config.gguf_model == "weights.gguf"
     assert config.unquantized_modules == ["proj_out"]
     assert config.get_name() == "gguf"
-    assert isinstance(config.get_vllm_quant_config(), _GGUFConfig)
 
 
-def test_get_diffusion_quant_config_builds_gguf_config():
-    config = get_diffusion_quant_config(
+def test_build_quant_config_builds_gguf_config():
+    config = build_quant_config(
         "gguf",
         gguf_model="repo/model-Q4_0.gguf",
         unquantized_modules=["norm_out"],
     )
 
-    assert isinstance(config, DiffusionGgufConfig)
+    assert isinstance(config, DiffusionGGUFConfig)
     assert config.gguf_model == "repo/model-Q4_0.gguf"
     assert config.unquantized_modules == ["norm_out"]
 
 
 def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
     linear = object.__new__(LinearBase)
-    method = _GGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
+    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
 
     assert isinstance(method, DiffusionGGUFLinearMethod)
 
 
 def test_gguf_vllm_config_respects_unquantized_modules():
     linear = object.__new__(LinearBase)
-    method = _GGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
+    method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
 
     assert isinstance(method, UnquantizedLinearMethod)
 
 
 def test_gguf_vllm_config_returns_none_for_non_linear_layers():
-    method = _GGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
+    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
 
     assert method is None
 
@@ -81,7 +76,7 @@ def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
 
 def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
     qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[]))
+    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
     layer = SimpleNamespace(
         qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
         qweight_type=SimpleNamespace(weight_type=qweight_type),
@@ -97,7 +92,7 @@ def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
 
 def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
     qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(_GGUFConfig(unquantized_modules=[]))
+    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
     qweight = torch.nn.Parameter(
         torch.tensor(
             [

tests/diffusion/quantization/test_int8_config.py

@@ -10,11 +10,11 @@ from pytest_mock import MockerFixture
 from torch.nn import Module, Parameter
 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
 
-from vllm_omni.diffusion.quantization import (
-    get_diffusion_quant_config,
-    get_vllm_quant_config_for_layers,
-)
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
 
 npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.")
 
@@ -23,36 +23,26 @@ cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason=
 
 def test_int8_config_creation():
     """Test that Int8 config can be created."""
-    config = get_diffusion_quant_config("int8")
+    config = build_quant_config("int8")
     assert config is not None
     assert config.get_name() == "int8"
 
 
-def test_vllm_config_extraction():
-    """Test that vLLM config can be extracted from diffusion config."""
-    diff_config = get_diffusion_quant_config("int8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
-    assert vllm_config is not None
-    assert vllm_config.activation_scheme == "dynamic"
-
-
 def test_none_quantization():
     """Test that None quantization returns None config."""
-    config = get_diffusion_quant_config(None)
+    config = build_quant_config(None)
     assert config is None
-    vllm_config = get_vllm_quant_config_for_layers(config)
-    assert vllm_config is None
 
 
 def test_invalid_quantization():
     """Test that invalid quantization method raises error."""
     with pytest.raises(ValueError, match="Unknown quantization method"):
-        get_diffusion_quant_config("invalid_method")
+        build_quant_config("invalid_method")
 
 
 def test_int8_config_with_custom_params():
     """Test Int8 config with custom parameters."""
-    config = get_diffusion_quant_config(
+    config = build_quant_config(
         "int8",
         activation_scheme="dynamic",
         ignored_layers=["proj_out"],
@@ -64,8 +54,6 @@ def test_int8_config_with_custom_params():
 
 def test_supported_methods():
     """Test that supported methods list is correct."""
-    from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
-
     assert "int8" in SUPPORTED_QUANTIZATION_METHODS
 
 
@@ -73,8 +61,8 @@ def test_quantization_integration():
     """Test end-to-end quantization flow through OmniDiffusionConfig."""
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    # Test with quantization string only
-    config = OmniDiffusionConfig(model="test", quantization="int8")
+    # Test with quantization_config string
+    config = OmniDiffusionConfig(model="test", quantization_config="int8")
     assert config.quantization_config is not None
     assert config.quantization_config.get_name() == "int8"
 
@@ -87,10 +75,6 @@ def test_quantization_integration():
     assert config2.quantization_config.get_name() == "int8"
     assert config2.quantization_config.activation_scheme == "dynamic"
 
-    # Test that vLLM config can be extracted
-    vllm_config = config.quantization_config.get_vllm_quant_config()
-    assert vllm_config is not None
-
 
 def test_quantization_dict_not_mutated():
     """Test that passing a dict to quantization_config doesn't mutate it."""
@@ -105,28 +89,24 @@ def test_quantization_dict_not_mutated():
     assert original_dict == dict_copy
 
 
-def test_quantization_conflicting_methods_warning(caplog):
-    """Test warning when quantization and quantization_config['method'] conflict."""
-    import logging
-
+def test_quantization_config_string_and_dict_equivalent():
+    """Test that string and dict forms produce equivalent configs."""
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-    with caplog.at_level(logging.WARNING):
-        config = OmniDiffusionConfig(
-            model="test",
-            quantization="int8",  # This should be overridden
-            quantization_config={"method": "int8", "activation_scheme": "dynamic"},
-        )
-    # No warning when methods match
-    assert config.quantization_config is not None
+    config_str = OmniDiffusionConfig(model="test", quantization_config="int8")
+    config_dict = OmniDiffusionConfig(
+        model="test",
+        quantization_config={"method": "int8", "activation_scheme": "dynamic"},
+    )
+    assert config_str.quantization_config.get_name() == config_dict.quantization_config.get_name()
+    assert config_str.quantization_config.activation_scheme == config_dict.quantization_config.activation_scheme
 
 
 def test_get_quant_method(mocker: MockerFixture):
     """Test for get_quant_method method for GPU"""
-    from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+    from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
 
-    diff_config = get_diffusion_quant_config("int8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+    config = build_quant_config("int8")
 
     def _fake_init(self, quant_config):
         pass
@@ -141,21 +121,20 @@ def test_get_quant_method(mocker: MockerFixture):
         patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True),
         patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False),
     ):
-        method = vllm_config.get_quant_method(layer, prefix)
+        method = config.get_quant_method(layer, prefix)
         assert isinstance(method, Int8OnlineLinearMethod)
 
     # Test skipping quantization for a layer
-    vllm_config.ignored_layers = [prefix]
-    method = vllm_config.get_quant_method(layer, prefix)
+    config.ignored_layers = [prefix]
+    method = config.get_quant_method(layer, prefix)
     assert isinstance(method, UnquantizedLinearMethod)
 
 
 def test_get_npu_quant_method():
     """Test for get_quant_method method for NPU"""
-    from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
 
-    diff_config = get_diffusion_quant_config("int8")
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+    config = build_quant_config("int8")
 
     layer = MagicMock(spec=LinearBase)
     prefix = "test_layer"
@@ -165,12 +144,12 @@ def test_get_npu_quant_method():
         patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False),
         patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True),
     ):
-        method = vllm_config.get_quant_method(layer, prefix)
+        method = config.get_quant_method(layer, prefix)
         assert isinstance(method, NPUInt8OnlineLinearMethod)
 
     # Test skipping quantization for a layer
-    vllm_config.ignored_layers = [prefix]
-    method = vllm_config.get_quant_method(layer, prefix)
+    config.ignored_layers = [prefix]
+    method = config.get_quant_method(layer, prefix)
     assert isinstance(method, UnquantizedLinearMethod)
 
 
@@ -189,12 +168,12 @@ class TestInt8LinearMethod:
     @pytest.fixture
     def patch_deps(self, mocker, mock_kernel):
         # mock init_int8_linear_kernel
-        mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
+        mocker.patch("vllm_omni.quantization.int8_config.init_int8_linear_kernel", return_value=mock_kernel)
         return mock_kernel
 
     def test_init(self, patch_deps, mock_quant_config):
         # test for Int8LinearMethod init
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod, init_int8_linear_kernel
+        from vllm_omni.quantization.int8_config import Int8LinearMethod, init_int8_linear_kernel
 
         method = Int8LinearMethod(mock_quant_config)
 
@@ -205,7 +184,7 @@ class TestInt8LinearMethod:
         assert method.int8_linear == patch_deps
 
     def test_process_weights_after_loading(self, patch_deps, mock_quant_config):
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+        from vllm_omni.quantization.int8_config import Int8LinearMethod
 
         method = Int8LinearMethod(mock_quant_config)
         layer = Module()
@@ -214,7 +193,7 @@ class TestInt8LinearMethod:
         patch_deps.process_weights_after_loading.assert_called_once_with(layer)
 
     def test_apply(self, patch_deps, mock_quant_config):
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+        from vllm_omni.quantization.int8_config import Int8LinearMethod
 
         method = Int8LinearMethod(mock_quant_config)
         layer = Module()
@@ -237,19 +216,19 @@ class TestInt8OnlineLinearMethod:
         # mock kernel
         mock_kernel = mocker.Mock()
         mock_kernel.layer_param_names = ("weight", "weight_scale", "input_scale", "input_zero_point", "azp_adj")
-        mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
-        mocker.patch("vllm_omni.diffusion.quantization.int8.replace_parameter")
+        mocker.patch("vllm_omni.quantization.int8_config.init_int8_linear_kernel", return_value=mock_kernel)
+        mocker.patch("vllm_omni.quantization.int8_config.replace_parameter")
 
         # mock scaled_int8_quant return value
         mock_qweight = torch.ones((128, 64), dtype=torch.int8)
         mock_scale = torch.randn(128)
         mock_quant = mocker.patch(
-            "vllm_omni.diffusion.quantization.int8.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
+            "vllm_omni.quantization.int8_config.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
         )
         return {"kernel": mock_kernel, "quant": mock_quant, "mock_qweight": mock_qweight, "mock_scale": mock_scale}
 
     def test_process_weights_after_loading(self, mock_deps, mock_quant_config):
-        from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
 
         method = Int8OnlineLinearMethod(mock_quant_config)
         layer = Module()
@@ -268,12 +247,12 @@ class TestNPUInt8LinearMethod:
     def mock_torch_npu(self, mocker):
         torch_npu = MagicMock()
 
-        mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu", return_value=torch_npu)
+        mocker.patch("vllm_omni.quantization.int8_config.torch_npu", return_value=torch_npu)
         mocker.patch(
-            "vllm_omni.diffusion.quantization.int8.torch_npu.npu_dynamic_quant",
+            "vllm_omni.quantization.int8_config.torch_npu.npu_dynamic_quant",
             return_value=(self.qweight_mock, self.scale_mock),
         )
-        mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu.npu_quant_matmul", return_value=self.out_mock)
+        mocker.patch("vllm_omni.quantization.int8_config.torch_npu.npu_quant_matmul", return_value=self.out_mock)
         return torch_npu
 
     @pytest.fixture
@@ -288,7 +267,7 @@ class TestNPUInt8LinearMethod:
         return layer
 
     def test_npu_int8_process_weights_after_loading(self, mock_layer, mock_quant_config, mock_torch_npu):
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
 
         method = NPUInt8LinearMethod(mock_quant_config)
         ori_weight_shape = mock_layer.weight.shape
@@ -299,7 +278,7 @@ class TestNPUInt8LinearMethod:
         assert mock_layer.weight.is_contiguous()
 
     def test_npu_int8_apply(self, mock_layer, mock_quant_config, mock_torch_npu):
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
 
         method = NPUInt8LinearMethod(mock_quant_config)
         x = torch.randn(1, 16, 64)
@@ -308,7 +287,7 @@ class TestNPUInt8LinearMethod:
         assert output.shape == (1, 16, 128)
 
     def test_npu_int8_online_process_weights(self, mock_layer, mock_quant_config, mock_torch_npu):
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
 
         method = NPUInt8OnlineLinearMethod(mock_quant_config)
         method.process_weights_after_loading(mock_layer)
@@ -320,9 +299,9 @@ class TestNPUInt8LinearMethod:
 @pytest.fixture
 def quant_config():
     """Shared quant config fixture for smoke tests."""
-    from vllm_omni.diffusion.quantization.int8 import Int8Config
+    from vllm_omni.quantization.int8_config import DiffusionInt8Config
 
-    return Int8Config(
+    return DiffusionInt8Config(
         is_checkpoint_int8_serialized=False,
         activation_scheme="dynamic",
     )
@@ -360,7 +339,7 @@ class TestNPUInt8Smoke:
 
     def test_real_npu_online_process_weights_after_loading(self, quant_config, real_layer):
         """Smoke test: full process_weights_after_loading with real torch_npu."""
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
 
         method = NPUInt8OnlineLinearMethod(quant_config)
 
@@ -375,7 +354,7 @@ class TestNPUInt8Smoke:
         """Smoke test: forward pass with real npu_quant_matmul."""
         import torch_npu
 
-        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+        from vllm_omni.quantization.int8_config import NPUInt8LinearMethod
 
         method = NPUInt8LinearMethod(quant_config)
 
@@ -425,7 +404,7 @@ class TestCudaInt8Smoke:
 
     def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
         """Smoke test: full process_weights_after_loading with real CUDA ops."""
-        from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+        from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
 
         method = Int8OnlineLinearMethod(quant_config)
 
@@ -439,7 +418,7 @@ class TestCudaInt8Smoke:
         """Smoke test: forward pass with real CUDA int8 kernel."""
         from vllm import _custom_ops as ops
 
-        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+        from vllm_omni.quantization.int8_config import Int8LinearMethod
 
         method = Int8LinearMethod(quant_config)
 

tests/diffusion/quantization/test_quantization_quality.py

@@ -0,0 +1,240 @@
+# 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.
+
+Developers adding a new quantization method should:
+1. Add their method + model to QUALITY_CONFIGS below
+2. Set a max_lpips threshold (use 0.15 for image, 0.20 for video as defaults)
+3. Run: pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
+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.
+
+Requirements:
+    pip install lpips
+
+Example — run only FP8 tests:
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "fp8"
+
+Example — run a specific model:
+    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
+"""
+
+from __future__ import annotations
+
+import gc
+from dataclasses import dataclass
+
+import numpy as np
+import pytest
+import torch
+
+from tests.utils import hardware_marks
+
+# ---------------------------------------------------------------------------
+# Configuration — add new quantization methods / models here
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+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
+    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
+
+
+# Add new quantization methods / models here.
+# Developers: copy a config, change quantization + max_lpips, run the test.
+QUALITY_CONFIGS = [
+    QualityTestConfig(
+        id="fp8_z_image",
+        model="Tongyi-MAI/Z-Image-Turbo",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table, morning light",
+        max_lpips=0.10,
+        num_inference_steps=20,
+    ),
+    QualityTestConfig(
+        id="fp8_flux",
+        model="black-forest-labs/FLUX.1-dev",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table, morning light",
+        max_lpips=0.20,
+        num_inference_steps=10,
+    ),
+    QualityTestConfig(
+        id="fp8_qwen_image",
+        model="Qwen/Qwen-Image",
+        quantization="fp8",
+        task="t2i",
+        prompt="a cup of coffee on a wooden table, morning light",
+        max_lpips=0.35,
+        seed=142,
+        num_inference_steps=20,
+    ),
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _generate_image(omni, config: QualityTestConfig):
+    """Generate a single image, return (PIL.Image, peak_mem_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(
+        device=current_omni_platform.device_type,
+    ).manual_seed(config.seed)
+    torch.cuda.reset_peak_memory_stats()
+
+    outputs = omni.generate(
+        {"prompt": config.prompt},
+        OmniDiffusionSamplingParams(
+            height=config.height,
+            width=config.width,
+            generator=generator,
+            num_inference_steps=config.num_inference_steps,
+        ),
+    )
+
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    first = outputs[0]
+    req_out = first.request_output[0] if hasattr(first, "request_output") else first
+    return req_out.images[0], peak_mem
+
+
+def _generate_video(omni, config: QualityTestConfig):
+    """Generate a video, return (np.ndarray [F,H,W,C], peak_mem_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(
+        device=current_omni_platform.device_type,
+    ).manual_seed(config.seed)
+    torch.cuda.reset_peak_memory_stats()
+
+    outputs = omni.generate(
+        {"prompt": config.prompt, "negative_prompt": ""},
+        OmniDiffusionSamplingParams(
+            height=config.height,
+            width=config.width,
+            generator=generator,
+            num_inference_steps=config.num_inference_steps,
+            num_frames=config.num_frames,
+        ),
+    )
+
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    first = outputs[0]
+    if hasattr(first, "request_output") and isinstance(first.request_output, list):
+        inner = first.request_output[0]
+        if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
+            frames = inner.images[0] if inner.images else None
+        else:
+            frames = inner
+    elif hasattr(first, "images") and first.images:
+        frames = first.images
+    else:
+        raise ValueError("Could not extract video frames from output.")
+
+    if isinstance(frames, torch.Tensor):
+        video = frames.detach().cpu()
+        if video.dim() == 5:
+            video = video[0]
+        if video.dim() == 4 and video.shape[0] in (3, 4):
+            video = video.permute(1, 2, 3, 0)
+        if video.is_floating_point():
+            video = video.clamp(-1, 1) * 0.5 + 0.5
+        return video.float().numpy(), peak_mem
+
+    return np.asarray(frames), peak_mem
+
+
+def _compute_lpips(baseline, quantized, task: str) -> float:
+    """Compute LPIPS between baseline and quantized outputs."""
+    from benchmarks.diffusion.quantization_quality import (
+        compute_lpips_images,
+        compute_lpips_video,
+    )
+
+    if task == "t2i":
+        return compute_lpips_images([baseline], [quantized])[0]
+    return compute_lpips_video(baseline, quantized)
+
+
+def _unload(omni):
+    del omni
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+_marks = hardware_marks(res={"cuda": "H100"})
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "config",
+    [pytest.param(c, id=c.id, marks=_marks) for c in QUALITY_CONFIGS],
+)
+def test_quantization_quality(config: QualityTestConfig):
+    """Validate that quantized output stays within LPIPS threshold of BF16."""
+    from vllm_omni.entrypoints.omni import Omni
+
+    generate_fn = _generate_video if config.task == "t2v" else _generate_image
+
+    # --- BF16 baseline ---
+    omni_bl = Omni(model=config.model)
+    baseline_out, bl_mem = generate_fn(omni_bl, config)
+    _unload(omni_bl)
+
+    # --- Quantized ---
+    omni_qt = Omni(model=config.model, quantization_config=config.quantization)
+    quant_out, qt_mem = generate_fn(omni_qt, config)
+    _unload(omni_qt)
+
+    # --- LPIPS ---
+    lpips_score = _compute_lpips(baseline_out, quant_out, config.task)
+
+    # --- 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"  LPIPS:         {lpips_score:.4f}  (threshold: {config.max_lpips})")
+    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}"
+    )

tests/e2e/offline_inference/compute_lpips.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Compute LPIPS between pairs of BF16 (baseline) and FP8 (quantized) images.
+
+Reads image pairs from a directory, computes LPIPS perceptual distance,
+and prints a Markdown results table.
+
+Usage:
+    python compute_lpips.py --image-dir ./test_quant_outputs
+
+Expects files named: <model>_bf16.png and <model>_fp8.png
+e.g. zimage_bf16.png / zimage_fp8.png, qwen_bf16.png / qwen_fp8.png
+
+Requirements:
+    pip install lpips Pillow torchvision
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+
+def compute_lpips(img_baseline, img_quantized, net="alex"):
+    """Compute LPIPS between two PIL images."""
+    import lpips
+    import torch
+    from torchvision import transforms
+
+    loss_fn = lpips.LPIPS(net=net).eval()
+    if torch.cuda.is_available():
+        loss_fn = loss_fn.cuda()
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize((256, 256)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+
+    t_bl = transform(img_baseline.convert("RGB")).unsqueeze(0)
+    t_qt = transform(img_quantized.convert("RGB")).unsqueeze(0)
+    if torch.cuda.is_available():
+        t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
+    with torch.no_grad():
+        score = loss_fn(t_bl, t_qt).item()
+    return score
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Compute LPIPS for BF16 vs FP8 image pairs.")
+    parser.add_argument(
+        "--image-dir", type=str, required=True, help="Directory containing *_bf16.png and *_fp8.png pairs."
+    )
+    parser.add_argument("--threshold", type=float, default=0.1, help="LPIPS threshold for PASS/FAIL (default: 0.1).")
+    parser.add_argument(
+        "--net", type=str, default="alex", choices=["alex", "vgg", "squeeze"], help="LPIPS backbone (default: alex)."
+    )
+    args = parser.parse_args()
+
+    from PIL import Image
+
+    image_dir = Path(args.image_dir)
+    if not image_dir.exists():
+        print(f"ERROR: directory {image_dir} does not exist")
+        sys.exit(1)
+
+    # Find all *_bf16.png files and look for matching *_fp8.png
+    bf16_files = sorted(image_dir.glob("*_bf16.png"))
+    if not bf16_files:
+        print(f"No *_bf16.png files found in {image_dir}")
+        sys.exit(0)
+
+    results = []
+    all_pass = True
+
+    for bf16_path in bf16_files:
+        model_name = bf16_path.stem.replace("_bf16", "")
+        fp8_path = image_dir / f"{model_name}_fp8.png"
+        if not fp8_path.exists():
+            print(f"  SKIP {model_name}: no matching {fp8_path.name}")
+            continue
+
+        img_bl = Image.open(bf16_path)
+        img_fp8 = Image.open(fp8_path)
+        score = compute_lpips(img_bl, img_fp8, net=args.net)
+
+        status = "PASS" if score < args.threshold else "FAIL"
+        if score >= args.threshold:
+            all_pass = False
+        results.append(
+            (model_name, score, status, f"{img_bl.width}x{img_bl.height}", f"{img_fp8.width}x{img_fp8.height}")
+        )
+
+    # Print results table
+    print("")
+    print("=" * 70)
+    print(f"  LPIPS Results (net={args.net}, threshold={args.threshold})")
+    print("=" * 70)
+    print(f"  {'Model':<20} {'LPIPS':>8} {'Status':>8}  {'BF16 size':>12} {'FP8 size':>12}")
+    print(f"  {'-' * 20} {'-' * 8} {'-' * 8}  {'-' * 12} {'-' * 12}")
+    for model_name, score, status, sz_bl, sz_fp8 in results:
+        print(f"  {model_name:<20} {score:>8.4f} {status:>8}  {sz_bl:>12} {sz_fp8:>12}")
+    print("")
+    print("  LPIPS < 0.01 = imperceptible")
+    print("  LPIPS < 0.05 = minor differences")
+    print("  LPIPS < 0.10 = noticeable but acceptable")
+    print("  LPIPS > 0.10 = clearly different")
+    print("=" * 70)
+
+    # Markdown table for PR
+    md_path = image_dir / "lpips_results.md"
+    with open(md_path, "w") as f:
+        f.write("## LPIPS Quality Benchmark (BF16 vs FP8)\n\n")
+        f.write(f"LPIPS backbone: `{args.net}` | threshold: `{args.threshold}`\n\n")
+        f.write("| Model | LPIPS | Status | BF16 Size | FP8 Size |\n")
+        f.write("|-------|-------|--------|-----------|----------|\n")
+        for model_name, score, status, sz_bl, sz_fp8 in results:
+            emoji = "✅" if status == "PASS" else "❌"
+            f.write(f"| {model_name} | {score:.4f} | {emoji} {status} | {sz_bl} | {sz_fp8} |\n")
+        f.write(f"\n> LPIPS < {args.threshold} = PASS\n")
+    print(f"  Markdown saved to: {md_path}")
+
+    if not all_pass:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

tests/e2e/offline_inference/run_quantization_e2e.sh

@@ -0,0 +1,200 @@
+#!/bin/bash
+# End-to-end test script for unified quantization framework (PR #1764).
+#
+# Tests FP8 quantization for:
+#   1. Z-Image-Turbo  (single-stage, ~20GB VRAM)
+#   2. Qwen-Image     (single-stage, ~20GB VRAM)
+#   3. FLUX.1-dev     (single-stage, ~25GB VRAM with fp8)
+#   4. BAGEL          (multi-stage LLM+DiT, ~55GB VRAM)
+#
+# Usage:
+#   bash tests/e2e/offline_inference/run_quantization_e2e.sh [--skip-flux] [--skip-bagel]
+#
+# Expected: all runs produce images without errors.
+# Key check for BAGEL: FP8 only applies to diffusion stage, NOT the LLM stage.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+OUTPUT_DIR="${REPO_ROOT}/test_quant_outputs"
+mkdir -p "$OUTPUT_DIR"
+
+SKIP_FLUX=false
+SKIP_BAGEL=false
+for arg in "$@"; do
+    case "$arg" in
+        --skip-flux) SKIP_FLUX=true ;;
+        --skip-bagel) SKIP_BAGEL=true ;;
+    esac
+done
+
+PASS=0
+FAIL=0
+SKIP=0
+
+run_test() {
+    local name="$1"
+    shift
+    echo ""
+    echo "============================================================"
+    echo "  TEST: $name"
+    echo "============================================================"
+    if "$@"; then
+        echo "  PASS: $name"
+        PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $name (exit code $?)"
+        FAIL=$((FAIL + 1))
+    fi
+}
+
+# ─── 1. Z-Image-Turbo ────────────────────────────────────────────────────────
+# Defaults from README: 1024x1024, 50 steps, cfg-scale 4.0, seed 42
+
+run_test "Z-Image-Turbo BF16 (baseline)" \
+    python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --prompt "a cup of coffee on the table" \
+        --seed 42 --num-inference-steps 50 \
+        --height 1024 --width 1024 \
+        --cfg-scale 4.0 \
+        --output "$OUTPUT_DIR/zimage_bf16.png"
+
+run_test "Z-Image-Turbo FP8" \
+    python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --prompt "a cup of coffee on the table" \
+        --seed 42 --num-inference-steps 50 \
+        --height 1024 --width 1024 \
+        --cfg-scale 4.0 \
+        --quantization fp8 \
+        --output "$OUTPUT_DIR/zimage_fp8.png"
+
+# ─── 2. Qwen-Image ───────────────────────────────────────────────────────────
+# Defaults from README: 1024x1024, 50 steps, cfg-scale 4.0, seed 142
+
+run_test "Qwen-Image BF16 (baseline)" \
+    python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
+        --model Qwen/Qwen-Image \
+        --prompt "a cup of coffee on the table" \
+        --seed 142 --num-inference-steps 50 \
+        --height 1024 --width 1024 \
+        --cfg-scale 4.0 \
+        --output "$OUTPUT_DIR/qwen_bf16.png"
+
+run_test "Qwen-Image FP8" \
+    python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
+        --model Qwen/Qwen-Image \
+        --prompt "a cup of coffee on the table" \
+        --seed 142 --num-inference-steps 50 \
+        --height 1024 --width 1024 \
+        --cfg-scale 4.0 \
+        --quantization fp8 \
+        --output "$OUTPUT_DIR/qwen_fp8.png"
+
+# ─── 3. FLUX.1-dev ───────────────────────────────────────────────────────────
+# Defaults from first test run: 1024x1024, 20 steps, guidance-scale 3.5
+
+if [ "$SKIP_FLUX" = true ]; then
+    echo ""
+    echo "  SKIP: FLUX.1-dev (--skip-flux)"
+    SKIP=$((SKIP + 1))
+    SKIP=$((SKIP + 1))
+else
+    run_test "FLUX.1-dev BF16 (baseline)" \
+        python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
+            --model black-forest-labs/FLUX.1-dev \
+            --prompt "a cup of coffee on the table" \
+            --seed 42 --num-inference-steps 20 \
+            --height 1024 --width 1024 \
+            --guidance-scale 3.5 --cfg-scale 1.0 \
+            --output "$OUTPUT_DIR/flux_bf16.png"
+
+    run_test "FLUX.1-dev FP8" \
+        python "$REPO_ROOT/examples/offline_inference/text_to_image/text_to_image.py" \
+            --model black-forest-labs/FLUX.1-dev \
+            --prompt "a cup of coffee on the table" \
+            --seed 42 --num-inference-steps 20 \
+            --height 1024 --width 1024 \
+            --guidance-scale 3.5 --cfg-scale 1.0 \
+            --quantization fp8 \
+            --output "$OUTPUT_DIR/flux_fp8.png"
+fi
+
+# ─── 4. BAGEL (multi-stage) ──────────────────────────────────────────────────
+# Defaults from README: 50 steps, cfg-text-scale 4.0, cfg-img-scale 1.5
+# BAGEL end2end.py saves images as output_1_stage_None_0.png in cwd.
+# We copy them to OUTPUT_DIR with proper names for LPIPS comparison.
+
+if [ "$SKIP_BAGEL" = true ]; then
+    echo ""
+    echo "  SKIP: BAGEL (--skip-bagel)"
+    SKIP=$((SKIP + 1))
+    SKIP=$((SKIP + 1))
+else
+    BAGEL_WORKDIR=$(mktemp -d)
+
+    run_test "BAGEL BF16 (baseline)" \
+        bash -c "cd '$BAGEL_WORKDIR' && python '$REPO_ROOT/examples/offline_inference/bagel/end2end.py' \
+            --model ByteDance-Seed/BAGEL-7B-MoT \
+            --modality text2img \
+            --prompts 'A cute cat' \
+            --steps 50 \
+            --cfg-text-scale 4.0 \
+            --cfg-img-scale 1.5 \
+            --seed 52 \
+            && cp -f output_1_stage_None_0.png '$OUTPUT_DIR/bagel_bf16.png'"
+
+    run_test "BAGEL FP8 (diffusion-only quantization)" \
+        bash -c "cd '$BAGEL_WORKDIR' && python '$REPO_ROOT/examples/offline_inference/bagel/end2end.py' \
+            --model ByteDance-Seed/BAGEL-7B-MoT \
+            --modality text2img \
+            --prompts 'A cute cat' \
+            --steps 50 \
+            --cfg-text-scale 4.0 \
+            --cfg-img-scale 1.5 \
+            --seed 52 \
+            --quantization fp8 \
+            && cp -f output_1_stage_None_0.png '$OUTPUT_DIR/bagel_fp8.png'"
+
+    rm -rf "$BAGEL_WORKDIR"
+fi
+
+# ─── 5. LPIPS quality comparison ─────────────────────────────────────────────
+# Compares each BF16/FP8 image pair using perceptual LPIPS distance.
+# Requires: pip install lpips torchvision
+
+echo ""
+echo "============================================================"
+echo "  LPIPS Quality Comparison"
+echo "============================================================"
+
+if python -c "import lpips" 2>/dev/null; then
+    run_test "LPIPS BF16-vs-FP8 quality check" \
+        python "$SCRIPT_DIR/compute_lpips.py" \
+            --image-dir "$OUTPUT_DIR" \
+            --threshold 0.1
+else
+    echo "  SKIP: lpips not installed (pip install lpips torchvision)"
+    SKIP=$((SKIP + 1))
+fi
+
+# ─── Summary ─────────────────────────────────────────────────────────────────
+
+echo ""
+echo "============================================================"
+echo "  SUMMARY"
+echo "============================================================"
+echo "  PASS: $PASS"
+echo "  FAIL: $FAIL"
+echo "  SKIP: $SKIP"
+echo "  Output images: $OUTPUT_DIR/"
+if [ -f "$OUTPUT_DIR/lpips_results.md" ]; then
+    echo "  LPIPS report: $OUTPUT_DIR/lpips_results.md"
+fi
+echo "============================================================"
+
+if [ "$FAIL" -gt 0 ]; then
+    exit 1
+fi

tests/e2e/offline_inference/test_quantization_fp8.py

@@ -0,0 +1,323 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end tests for the unified quantization framework (PR #1764).
+
+Validates FP8 quantization works correctly for all supported model types:
+  - Single-stage diffusion models (FLUX.1-dev, Qwen-Image, Z-Image-Turbo)
+  - Multi-stage models (BAGEL: LLM + Diffusion)
+
+Tests verify:
+  1. FP8 quantization produces valid images
+  2. Memory usage is lower than BF16 baseline
+  3. Multi-stage models only quantize the diffusion stage (not the LLM stage)
+
+Usage:
+    # Run all FP8 quantization tests
+    pytest tests/e2e/offline_inference/test_quantization_fp8.py -v
+
+    # Run single-stage tests only (faster, needs ~25GB VRAM)
+    pytest tests/e2e/offline_inference/test_quantization_fp8.py -v -k "single_stage"
+
+    # Run BAGEL multi-stage test (needs ~55GB VRAM, H100 recommended)
+    pytest tests/e2e/offline_inference/test_quantization_fp8.py -v -k "bagel"
+
+    # Run FLUX test only
+    pytest tests/e2e/offline_inference/test_quantization_fp8.py -v -k "flux"
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+import torch
+
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+# ─── helpers ──────────────────────────────────────────────────────────────────
+
+
+def _generate_single_stage_image(
+    model: str,
+    quantization: str | None = None,
+    height: int = 256,
+    width: int = 256,
+    num_inference_steps: int = 2,
+    seed: int = 42,
+    **extra_omni_kwargs: Any,
+) -> tuple[list, float]:
+    """Generate an image with a single-stage diffusion model.
+
+    Returns (images, peak_memory_gib).
+    """
+    omni_kwargs: dict[str, Any] = {"model": model, **extra_omni_kwargs}
+    if quantization:
+        omni_kwargs["quantization"] = quantization
+
+    omni = Omni(**omni_kwargs)
+    try:
+        torch.cuda.reset_peak_memory_stats()
+
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+        outputs = omni.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=0.0,
+                generator=generator,
+            ),
+        )
+
+        peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+        first_output = outputs[0]
+        assert first_output.final_output_type == "image"
+        assert hasattr(first_output, "request_output") and first_output.request_output
+        req_out = first_output.request_output[0]
+        assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+        images = req_out.images
+        assert len(images) >= 1
+        assert images[0].width == width
+        assert images[0].height == height
+
+        return images, peak_mem
+    finally:
+        omni.close()
+
+
+def _generate_bagel_image(
+    quantization_config: str | None = None,
+    num_inference_steps: int = 15,
+) -> tuple[Any, float]:
+    """Generate an image with BAGEL (multi-stage: LLM + Diffusion).
+
+    Returns (generated_image, peak_memory_gib).
+    """
+    config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+    omni_kwargs: dict[str, Any] = {
+        "model": "ByteDance-Seed/BAGEL-7B-MoT",
+        "stage_configs_path": config_path,
+        "stage_init_timeout": 300,
+    }
+    if quantization_config:
+        omni_kwargs["quantization_config"] = quantization_config
+
+    omni = Omni(**omni_kwargs)
+    try:
+        torch.cuda.reset_peak_memory_stats()
+
+        params_list = omni.default_sampling_params_list
+        params_list[0].max_tokens = 1  # type: ignore
+        if len(params_list) > 1:
+            params_list[1].num_inference_steps = num_inference_steps  # type: ignore
+            params_list[1].extra_args = {  # type: ignore
+                "cfg_text_scale": 4.0,
+                "cfg_img_scale": 1.5,
+            }
+
+        prompt = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|im_end|>"
+        omni_outputs = list(
+            omni.generate(
+                prompts=[{"prompt": prompt, "modalities": ["image"]}],
+                sampling_params_list=params_list,
+            )
+        )
+
+        peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+        # Extract image
+        generated_image = None
+        for req_output in omni_outputs:
+            if images := getattr(req_output, "images", None):
+                generated_image = images[0]
+                break
+            if hasattr(req_output, "request_output") and req_output.request_output:
+                for stage_out in req_output.request_output:
+                    if hasattr(stage_out, "images") and stage_out.images:
+                        generated_image = stage_out.images[0]
+                        break
+                if generated_image:
+                    break
+
+        assert generated_image is not None, "No images generated from BAGEL"
+        assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
+
+        # Check LLM stage output — should have finish_reason=stop (not length)
+        for req_output in omni_outputs:
+            if hasattr(req_output, "request_output"):
+                for stage_out in req_output.request_output:
+                    if hasattr(stage_out, "outputs"):
+                        for comp_out in stage_out.outputs:
+                            if hasattr(comp_out, "finish_reason"):
+                                assert comp_out.finish_reason == "stop", (
+                                    f"LLM stage finish_reason={comp_out.finish_reason}, "
+                                    f"text={comp_out.text!r}. "
+                                    "FP8 may have leaked to the LLM stage."
+                                )
+
+        return generated_image, peak_mem
+    finally:
+        omni.close()
+
+
+# ─── Single-stage diffusion model tests ──────────────────────────────────────
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_single_stage_zimage_fp8():
+    """Z-Image-Turbo with FP8 generates valid images."""
+    images, _ = _generate_single_stage_image(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        quantization="fp8",
+    )
+    assert len(images) >= 1
+    images[0].save("test_zimage_fp8.png")
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_single_stage_zimage_fp8_uses_less_memory():
+    """FP8 should use less peak memory than BF16 for Z-Image-Turbo."""
+    _, mem_bf16 = _generate_single_stage_image(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        quantization=None,
+    )
+    torch.cuda.empty_cache()
+
+    _, mem_fp8 = _generate_single_stage_image(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        quantization="fp8",
+    )
+
+    print(f"Z-Image BF16 peak memory: {mem_bf16:.2f} GiB")
+    print(f"Z-Image FP8 peak memory:  {mem_fp8:.2f} GiB")
+    assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_single_stage_qwen_image_fp8():
+    """Qwen-Image (random weights) with FP8 generates valid images."""
+    model = "riverclouds/qwen_image_random"
+    if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
+        pytest.skip("qwen_image_random not available on this platform")
+
+    images, _ = _generate_single_stage_image(
+        model=model,
+        quantization="fp8",
+    )
+    assert len(images) >= 1
+    images[0].save("test_qwen_image_fp8.png")
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_single_stage_flux_fp8():
+    """FLUX.1-dev with FP8 generates valid images."""
+    images, _ = _generate_single_stage_image(
+        model="black-forest-labs/FLUX.1-dev",
+        quantization="fp8",
+        height=512,
+        width=512,
+        num_inference_steps=4,
+    )
+    assert len(images) >= 1
+    images[0].save("test_flux_fp8.png")
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_single_stage_flux_fp8_uses_less_memory():
+    """FP8 should use less peak memory than BF16 for FLUX.1-dev."""
+    _, mem_bf16 = _generate_single_stage_image(
+        model="black-forest-labs/FLUX.1-dev",
+        quantization=None,
+        height=512,
+        width=512,
+        num_inference_steps=4,
+    )
+    torch.cuda.empty_cache()
+
+    _, mem_fp8 = _generate_single_stage_image(
+        model="black-forest-labs/FLUX.1-dev",
+        quantization="fp8",
+        height=512,
+        width=512,
+        num_inference_steps=4,
+    )
+
+    print(f"FLUX BF16 peak memory: {mem_bf16:.2f} GiB")
+    print(f"FLUX FP8 peak memory:  {mem_fp8:.2f} GiB")
+    assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
+
+
+# ─── Multi-stage model tests (BAGEL) ─────────────────────────────────────────
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_bagel_fp8_generates_image():
+    """BAGEL with FP8 quantization_config generates a valid image.
+
+    FP8 should only apply to the diffusion stage (Stage-1), not the
+    LLM stage (Stage-0). We verify this by checking:
+      1. Image is generated successfully
+      2. LLM stage finish_reason is 'stop' (not 'length' from garbled output)
+    """
+    image, _ = _generate_bagel_image(quantization_config="fp8")
+    image.save("test_bagel_fp8.png")
+
+
+@hardware_test(res={"cuda": "H100"})
+def test_bagel_bf16_generates_image():
+    """BAGEL without quantization generates a valid image (baseline)."""
+    image, _ = _generate_bagel_image(quantization_config=None)
+    image.save("test_bagel_bf16.png")
+
+
+# ─── Quantization config routing tests ────────────────────────────────────────
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_quantization_key_maps_to_quantization_config():
+    """The old 'quantization' kwarg should map to 'quantization_config'
+    in OmniDiffusionConfig.from_kwargs for backwards compatibility."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig.from_kwargs(model="test", quantization="fp8")
+    assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "fp8"
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_quantization_config_key_takes_priority():
+    """When both 'quantization' and 'quantization_config' are set,
+    'quantization_config' takes priority."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig.from_kwargs(
+        model="test",
+        quantization="fp8",
+        quantization_config={"method": "fp8", "activation_scheme": "static"},
+    )
+    assert config.quantization_config is not None
+    assert config.quantization_config.activation_scheme == "static"
+
+
+@hardware_test(res={"cuda": "L4"})
+def test_single_stage_quantization_config_key():
+    """Single-stage model using quantization_config (dict) key generates images."""
+    images, _ = _generate_single_stage_image(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        quantization=None,
+        quantization_config="fp8",
+    )
+    assert len(images) >= 1

vllm_omni/config/model.py

@@ -6,12 +6,57 @@ from vllm.config import ModelConfig
 from vllm.config.utils import config
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_text_config
+from vllm.transformers_utils.model_arch_config_convertor import (
+    ModelArchConfigConvertorBase,
+)
 
 import vllm_omni.model_executor.models as me_models
 
 logger = init_logger(__name__)
 
 
+class OmniModelArchConfigConvertor(ModelArchConfigConvertorBase):
+    """Config convertor for Omni multi-stage models.
+
+    Pre-quantized checkpoints (e.g. modelopt FP8) store quantization
+    config in a stage-specific sub-config (e.g.
+    thinker_config.text_config.quantization_config) with correct relative
+    prefixes.  The legacy hf_quant_config.json sits at the top level with
+    "thinker."-prefixed names that don't match vllm-omni's module names.
+
+    This convertor accepts an optional *stage_config_name* so that only
+    the relevant stage's quantization config is surfaced.
+    """
+
+    def __init__(
+        self,
+        hf_config,
+        hf_text_config,
+        stage_config_name: str | None = None,
+    ):
+        super().__init__(hf_config, hf_text_config)
+        self.stage_config_name = stage_config_name
+
+    def get_quantization_config(self):
+        # When a stage_config_name is set, look for quantization config
+        # in that stage's text_config first (has correct relative prefixes).
+        if self.stage_config_name is not None:
+            stage_cfg = getattr(self.hf_config, self.stage_config_name, None)
+            if stage_cfg is not None:
+                text_cfg = getattr(stage_cfg, "text_config", None)
+                if text_cfg is not None:
+                    quant_cfg = self._normalize_quantization_config(text_cfg)
+                    if quant_cfg is not None:
+                        return quant_cfg
+
+            # For non-thinker stages (talker, code2wav) whose text_config
+            # has no quantization_config, return None so quantization is
+            # not applied to stages that were not quantized.
+            return None
+
+        return super().get_quantization_config()
+
+
 @config(config=ConfigDict(arbitrary_types_allowed=True))
 class OmniModelConfig(ModelConfig):
     """Configuration for Omni models, extending the base ModelConfig.
@@ -87,6 +132,20 @@ class OmniModelConfig(ModelConfig):
                 return override
         return super().embedding_size
 
+    def get_model_arch_config(self):
+        # For multi-stage omni models, use a stage-aware convertor so that
+        # only the correct stage's quantization config is surfaced.
+        # Without this, a pre-quantized thinker checkpoint would also
+        # apply quantization to the talker/code2wav stages.
+        if self.hf_config_name is not None:
+            convertor = OmniModelArchConfigConvertor(
+                self.hf_config,
+                self.hf_text_config,
+                stage_config_name=self.hf_config_name,
+            )
+            return convertor.convert()
+        return super().get_model_arch_config()
+
     def draw_hf_text_config(self):
         # transformers' get_text_config method is used to get the text config from thinker_config.
         # to handle the case that each model stage has their own text config,

vllm_omni/diffusion/data.py

@@ -13,18 +13,16 @@ from pydantic import model_validator
 from typing_extensions import Self
 from vllm.config.utils import config
 from vllm.logger import init_logger
-
-from vllm_omni.diffusion.quantization import (
-    DiffusionQuantizationConfig,
-    get_diffusion_quant_config,
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
 )
+
 from vllm_omni.diffusion.utils.network_utils import is_port_available
+from vllm_omni.quantization import build_quant_config
 
 if TYPE_CHECKING:
     from vllm.config import ProfilerConfig
 
-    from vllm_omni.diffusion.quantization import DiffusionQuantizationConfig
-
 # Import after TYPE_CHECKING to avoid circular imports at runtime
 # The actual import is deferred to __post_init__ to avoid import order issues
 
@@ -483,10 +481,10 @@ class OmniDiffusionConfig:
     # Model-specific function for collecting CFG KV caches (set at runtime)
     cfg_kv_collect_func: Any | None = None
 
-    # Quantization settings
-    # Supported methods: "fp8" (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs)
-    quantization: str | None = None
-    quantization_config: "DiffusionQuantizationConfig | dict[str, Any] | None" = None
+    # Quantization: str method name, dict config, QuantizationConfig, or None.
+    # str is resolved to {"method": <str>} internally.
+    # Per-component: {"transformer": {"method": "fp8"}, "vae": None}
+    quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
 
     # Diffusion pipeline Profiling config
     enable_diffusion_pipeline_profiler: bool = False
@@ -597,34 +595,17 @@ class OmniDiffusionConfig:
             # If it's neither dict nor DiffusionCacheConfig, convert to empty config
             self.cache_config = DiffusionCacheConfig()
 
-        # Convert quantization config (deferred import to avoid circular imports)
-        if self.quantization is not None or self.quantization_config is not None:
-            from vllm_omni.diffusion.quantization import (
-                DiffusionQuantizationConfig,
-            )
-
-            # Handle dict or DictConfig (from OmegaConf) - use Mapping for broader compatibility
-            if isinstance(self.quantization_config, Mapping):
-                # Convert DictConfig to dict if needed (OmegaConf compatibility)
-                config_dict = dict(self.quantization_config)
-                # Use get() instead of pop() to avoid mutating original dict
-                quant_method = config_dict.get("method", self.quantization)
-                # Filter out "method" key for kwargs
-                quant_kwargs = {k: v for k, v in config_dict.items() if k != "method"}
-
-                # Validate conflicting methods
-                if self.quantization is not None and quant_method is not None and quant_method != self.quantization:
-                    logger.warning(
-                        f"Conflicting quantization methods: quantization={self.quantization!r}, "
-                        f"quantization_config['method']={quant_method!r}. Using quantization_config['method']."
-                    )
-
-                self.quantization_config = get_diffusion_quant_config(quant_method, **quant_kwargs)
-            elif self.quantization_config is None and self.quantization is not None:
-                self.quantization_config = get_diffusion_quant_config(self.quantization)
-            elif not isinstance(self.quantization_config, DiffusionQuantizationConfig):
+        # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.
+        if self.quantization_config is not None:
+            if isinstance(self.quantization_config, QuantizationConfig):
+                pass  # Already built
+            elif isinstance(self.quantization_config, str):
+                self.quantization_config = build_quant_config(self.quantization_config)
+            elif isinstance(self.quantization_config, Mapping):
+                self.quantization_config = build_quant_config(dict(self.quantization_config))
+            else:
                 raise TypeError(
-                    f"quantization_config must be a DiffusionQuantizationConfig, dict, or None, "
+                    f"quantization_config must be str, dict, QuantizationConfig, or None, "
                     f"got {type(self.quantization_config)!r}"
                 )
 
@@ -646,6 +627,13 @@ class OmniDiffusionConfig:
                 kwargs["lora_scale"] = kwargs["static_lora_scale"]
             kwargs.pop("static_lora_scale", None)
 
+        # Backwards-compatibility: map "quantization" to "quantization_config"
+        # so callers using the old field name still work.
+        if "quantization" in kwargs and "quantization_config" not in kwargs:
+            kwargs["quantization_config"] = kwargs.pop("quantization")
+        else:
+            kwargs.pop("quantization", None)
+
         # Check environment variable as fallback for cache_backend
         # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND for backwards compatibility
         if "cache_backend" not in kwargs:

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -260,7 +260,7 @@ class DiffusersPipelineLoader:
     ) -> nn.Module:
         """Load a model with the given configurations."""
         # CPU offload + FP8: load weights on device for FP8 quantization
-        if load_device == "cpu" and od_config.quantization and od_config.quantization.lower() != "none":
+        if load_device == "cpu" and od_config.quantization_config is not None:
             load_device = device.type
             logger.info(f"Quantization enabled with CPU offload, using {load_device} for weight loading")
 
@@ -342,28 +342,35 @@ class DiffusersPipelineLoader:
         quant_config = od_config.quantization_config
         if quant_config is None:
             return False
-        # Fast path: mapping-style config (e.g., DictConfig)
-        if isinstance(quant_config, dict):
-            method = str(quant_config.get("method", "")).lower()
-            if method != "gguf":
-                return False
-            gguf_model = quant_config.get("gguf_model")
-            if not gguf_model:
-                raise ValueError("GGUF quantization requires quantization_config.gguf_model")
+
+        # New API: DiffusionGGUFConfig (QuantizationConfig subclass)
+        from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
+
+        if isinstance(quant_config, DiffusionGGUFConfig):
+            if quant_config.gguf_model is None:
+                raise ValueError("GGUF quantization requires gguf_model")
             return True
 
-        # Normal path: DiffusionQuantizationConfig
-        if not hasattr(quant_config, "get_name"):
-            # Fallback: if it carries gguf_model, treat as GGUF
-            gguf_model = getattr(quant_config, "gguf_model", None)
-            return bool(gguf_model)
-        is_gguf = quant_config.get_name() == "gguf"
-        if not is_gguf:
+        # Dict-style config: {"method": "gguf", "gguf_model": "..."}
+        if isinstance(quant_config, dict):
+            if quant_config.get("method") == "gguf":
+                if not quant_config.get("gguf_model"):
+                    raise ValueError("GGUF quantization requires gguf_model")
+                return True
             return False
-        gguf_model = getattr(quant_config, "gguf_model", None)
-        if gguf_model is None:
-            raise ValueError("GGUF quantization requires quantization_config.gguf_model")
-        return True
+
+        # Check by name for any config that reports as "gguf"
+        if hasattr(quant_config, "get_name") and quant_config.get_name() == "gguf":
+            gguf_model = getattr(quant_config, "gguf_model", None)
+            if gguf_model is None:
+                raise ValueError("GGUF quantization requires gguf_model")
+            return True
+
+        # Fallback: object with gguf_model attribute but no get_name
+        if hasattr(quant_config, "gguf_model") and quant_config.gguf_model:
+            return True
+
+        return False
 
     def _is_transformer_source(self, source: "ComponentSource") -> bool:
         if source.subfolder == "transformer":

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -28,6 +28,9 @@ from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
     RowParallelLinear,
 )
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+)
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.transformers_utils.configs.bagel import BagelConfig
@@ -68,16 +71,27 @@ def patchify(imgs, p):
 
 
 class MLPconnector(nn.Module):
-    def __init__(self, input_dim, output_dim, activation="gelu_pytorch_tanh"):
+    def __init__(
+        self,
+        input_dim,
+        output_dim,
+        activation="gelu_pytorch_tanh",
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.fc1 = ColumnParallelLinear(input_dim, output_dim, bias=True, gather_output=False)
+        self.fc1 = ColumnParallelLinear(
+            input_dim, output_dim, bias=True, gather_output=False, quant_config=quant_config, prefix=f"{prefix}.fc1"
+        )
         if activation == "gelu":
             self.act = nn.GELU()
         elif activation == "gelu_pytorch_tanh":
             self.act = nn.GELU(approximate="tanh")
         else:
             self.act = nn.ReLU()
-        self.fc2 = RowParallelLinear(output_dim, output_dim, bias=True, input_is_parallel=True)
+        self.fc2 = RowParallelLinear(
+            output_dim, output_dim, bias=True, input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.fc2"
+        )
 
     def forward(self, x):
         x_parallel, _ = self.fc1(x)
@@ -139,6 +153,8 @@ class BagelMLP(nn.Module):
         hidden_size: int,
         intermediate_size: int,
         hidden_act: str = "silu",
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.gate_proj = ColumnParallelLinear(
@@ -146,18 +162,24 @@ class BagelMLP(nn.Module):
             intermediate_size,
             bias=False,
             gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_proj",
         )
         self.up_proj = ColumnParallelLinear(
             hidden_size,
             intermediate_size,
             bias=False,
             gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.up_proj",
         )
         self.down_proj = RowParallelLinear(
             intermediate_size,
             hidden_size,
             input_is_parallel=True,
             bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
         )
         if hidden_act != "silu":
             raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.")
@@ -273,7 +295,14 @@ class PackedAttentionMoT(nn.Module):
       - qkv_proj_moe_gen : stacks q_proj_moe_gen + k_proj_moe_gen + v_proj_moe_gen (gen vae)
     """
 
-    def __init__(self, config, layer_idx: int | None = None, parallel_config: DiffusionParallelConfig | None = None):
+    def __init__(
+        self,
+        config,
+        layer_idx: int | None = None,
+        parallel_config: DiffusionParallelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
         self.layer_idx = layer_idx
         self.hidden_size = config.hidden_size
@@ -300,11 +329,15 @@ class PackedAttentionMoT(nn.Module):
             self.total_num_heads,
             self.total_num_kv_heads,
             bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
         )
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
             self.hidden_size,
             bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
         )
 
         # Generation mode MoE projections (stacked q/k/v)
@@ -314,11 +347,15 @@ class PackedAttentionMoT(nn.Module):
             self.total_num_heads,
             self.total_num_kv_heads,
             bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj_moe_gen",
         )
         self.o_proj_moe_gen = RowParallelLinear(
             self.total_num_heads * self.head_dim,
             self.hidden_size,
             bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj_moe_gen",
         )
 
         # QK normalization
@@ -607,14 +644,30 @@ class Qwen2MoTDecoderLayer(nn.Module):
         layer_idx: int | None = None,
         attn_module: type[nn.Module] | None = PackedAttentionMoT,
         parallel_config: DiffusionParallelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.hidden_size = config.hidden_size
 
-        self.self_attn = attn_module(config, layer_idx, parallel_config=parallel_config)
+        self.self_attn = attn_module(
+            config, layer_idx, parallel_config=parallel_config, quant_config=quant_config, prefix=f"{prefix}.self_attn"
+        )
 
-        self.mlp = BagelMLP(config.hidden_size, config.intermediate_size, config.hidden_act)
-        self.mlp_moe_gen = BagelMLP(config.hidden_size, config.intermediate_size, config.hidden_act)
+        self.mlp = BagelMLP(
+            config.hidden_size,
+            config.intermediate_size,
+            config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+        self.mlp_moe_gen = BagelMLP(
+            config.hidden_size,
+            config.intermediate_size,
+            config.hidden_act,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp_moe_gen",
+        )
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.input_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -693,7 +746,13 @@ class Qwen2MoTDecoderLayer(nn.Module):
 
 
 class Qwen2MoTModel(Qwen2PreTrainedModel):
-    def __init__(self, config, parallel_config: DiffusionParallelConfig | None = None):
+    def __init__(
+        self,
+        config,
+        parallel_config: DiffusionParallelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
@@ -702,7 +761,14 @@ class Qwen2MoTModel(Qwen2PreTrainedModel):
         self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
         self.layers = nn.ModuleList(
             [
-                Qwen2MoTDecoderLayer(config, layer_idx, attn_module=PackedAttentionMoT, parallel_config=parallel_config)
+                Qwen2MoTDecoderLayer(
+                    config,
+                    layer_idx,
+                    attn_module=PackedAttentionMoT,
+                    parallel_config=parallel_config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{layer_idx}",
+                )
                 for layer_idx in range(config.num_hidden_layers)
             ]
         )
@@ -784,9 +850,17 @@ class Qwen2MoTModel(Qwen2PreTrainedModel):
 class Qwen2MoTForCausalLM(Qwen2PreTrainedModel):
     _tied_weights_keys = ["lm_head.weight"]
 
-    def __init__(self, config, parallel_config: DiffusionParallelConfig | None = None):
+    def __init__(
+        self,
+        config,
+        parallel_config: DiffusionParallelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__(config)
-        self.model = Qwen2MoTModel(config, parallel_config=parallel_config)
+        self.model = Qwen2MoTModel(
+            config, parallel_config=parallel_config, quant_config=quant_config, prefix=f"{prefix}.model"
+        )
         self.vocab_size = config.vocab_size
 
         # Initialize weights and apply final processing
@@ -1001,6 +1075,8 @@ class Bagel(nn.Module):
         vit_model,
         config: BagelConfig,
         parallel_config: DiffusionParallelConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.language_model = language_model
@@ -1026,7 +1102,13 @@ class Bagel(nn.Module):
             self.vit_patch_size = config.vit_config.patch_size
             self.vit_max_num_patch_per_side = config.vit_max_num_patch_per_side
             self.vit_hidden_size = config.vit_config.hidden_size
-            self.connector = MLPconnector(self.vit_hidden_size, self.hidden_size, config.connector_act)
+            self.connector = MLPconnector(
+                self.vit_hidden_size,
+                self.hidden_size,
+                config.connector_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.connector",
+            )
             self.vit_pos_embed = PositionEmbedding(self.vit_max_num_patch_per_side, self.hidden_size)
 
         self.get_flattened_position_ids = get_flattened_position_ids_extrapolate

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -225,7 +225,14 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         )
 
         parallel_config = od_config.parallel_config if od_config else None
-        self.language_model = Qwen2MoTForCausalLM(llm_config, parallel_config=parallel_config)
+        quant_config = od_config.quantization_config
+        # Bagel uses explicit prefixes ("bagel.language_model", "bagel") because
+        # its model structure nests components under a top-level "bagel" module,
+        # unlike other pipelines where the transformer is the root module.
+        # This ensures ComponentQuantizationConfig prefix matching works correctly.
+        self.language_model = Qwen2MoTForCausalLM(
+            llm_config, parallel_config=parallel_config, quant_config=quant_config, prefix="bagel.language_model"
+        )
         ae_params: AutoEncoderParams = default_ae_params()
         self.vae = AutoEncoder(ae_params)
 
@@ -233,6 +240,8 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             language_model=self.language_model,
             vit_model=self.vit_model,
             parallel_config=parallel_config,
+            quant_config=quant_config,
+            prefix="bagel",
             config=BagelConfig(
                 llm_config=llm_config,
                 vae_config=vae_cfg,
@@ -257,7 +266,12 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             )
         ]
 
-        self.to(self.device)
+        # When quantization is enabled, vLLM linear layers live on meta
+        # device until the weight loader materializes them. Calling
+        # .to(device) would fail on those meta tensors, so we skip it
+        # entirely and let the weight loader handle device placement.
+        if quant_config is None:
+            self.to(self.device)
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )

vllm_omni/diffusion/models/flux/pipeline_flux.py

@@ -29,7 +29,6 @@ from vllm_omni.diffusion.models.flux import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin
 from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
@@ -107,8 +106,7 @@ class FluxPipeline(nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipe
         self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
         )
-        quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
-        self.transformer = FluxTransformer2DModel(od_config=od_config, quant_config=quant_config)
+        self.transformer = FluxTransformer2DModel(od_config=od_config, quant_config=od_config.quantization_config)
 
         self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
         self.tokenizer_2 = T5TokenizerFast.from_pretrained(

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -44,7 +44,6 @@ from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
 )
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -232,8 +231,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         ).to(self._execution_device)
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel)
-        quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
-        self.transformer = Flux2Transformer2DModel(quant_config=quant_config, **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/hunyuan_image_3/pipeline_hunyuan_image_3.py

@@ -20,7 +20,6 @@ from vllm.transformers_utils.config import get_config
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 from .autoencoder import AutoencoderKLConv3D
@@ -78,7 +77,7 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin):
                 fall_back_to_pt=True,
             )
         ]
-        quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
+        quant_config = od_config.quantization_config
         self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config)
         self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae)
         self._pipeline = None

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -33,7 +33,6 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
 
@@ -279,9 +278,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
             model, subfolder="vae", local_files_only=local_files_only
         ).to(self.device)
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
-        quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
         self.transformer = QwenImageTransformer2DModel(
-            od_config=od_config, quant_config=quant_config, **transformer_kwargs
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
         )
 
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -38,7 +38,6 @@ from vllm_omni.diffusion.models.z_image.z_image_transformer import (
     ZImageTransformer2DModel,
 )
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
@@ -175,9 +174,7 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         self.vae = DistributedAutoencoderKL.from_pretrained(
             model, subfolder="vae", local_files_only=local_files_only
         ).to(self._execution_device)
-        # Get vLLM quantization config for linear layers
-        quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
-        self.transformer = ZImageTransformer2DModel(quant_config=quant_config)
+        self.transformer = ZImageTransformer2DModel(quant_config=od_config.quantization_config)
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
         # Note: Context parallelism is applied centrally in registry.initialize_model()

vllm_omni/diffusion/quantization/__init__.py

@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Quantization support for diffusion models.
-
-This module provides a unified interface for quantizing diffusion transformers
-using various methods (FP8, etc.). It wraps vLLM's quantization infrastructure
-while allowing diffusion-model-specific defaults and optimizations.
-
-Example usage:
-    from vllm_omni.diffusion.quantization import (
-        get_diffusion_quant_config,
-        get_vllm_quant_config_for_layers,
-    )
-
-    # Create FP8 config for diffusion model
-    diff_config = get_diffusion_quant_config("fp8")
-
-    # Get vLLM config to pass to linear layers
-    vllm_config = get_vllm_quant_config_for_layers(diff_config)
-
-    # Use in model initialization
-    linear_layer = QKVParallelLinear(..., quant_config=vllm_config)
-"""
-
-from typing import TYPE_CHECKING
-
-from vllm.logger import init_logger
-
-from .base import DiffusionQuantizationConfig
-from .fp8 import DiffusionFp8Config
-from .gguf import DiffusionGgufConfig
-from .int8 import DiffusionInt8Config
-
-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization.base_config import (
-        QuantizationConfig,
-    )
-
-logger = init_logger(__name__)
-
-# Registry of supported quantization methods
-# To add a new method, create a new config class and register it here
-_QUANT_CONFIG_REGISTRY: dict[str, type[DiffusionQuantizationConfig]] = {
-    "fp8": DiffusionFp8Config,
-    "int8": DiffusionInt8Config,
-    "gguf": DiffusionGgufConfig,
-}
-
-SUPPORTED_QUANTIZATION_METHODS = list(_QUANT_CONFIG_REGISTRY.keys())
-
-
-def get_diffusion_quant_config(
-    quantization: str | None,
-    **kwargs,
-) -> DiffusionQuantizationConfig | None:
-    """Factory function to create quantization config for diffusion models.
-
-    Args:
-        quantization: Quantization method name ("fp8", etc.) or None to disable
-        **kwargs: Method-specific parameters passed to the config constructor
-
-    Returns:
-        DiffusionQuantizationConfig instance or None if quantization is disabled
-
-    Raises:
-        ValueError: If the quantization method is not supported
-
-    Example:
-        # Default FP8 with dynamic activation scaling
-        config = get_diffusion_quant_config("fp8")
-
-        # FP8 with custom parameters
-        config = get_diffusion_quant_config(
-            "fp8",
-            activation_scheme="static",
-            ignored_layers=["proj_out"],
-        )
-    """
-    if quantization is None or quantization.lower() == "none":
-        return None
-
-    quantization = quantization.lower()
-    if quantization not in _QUANT_CONFIG_REGISTRY:
-        raise ValueError(
-            f"Unknown quantization method: {quantization!r}. Supported methods: {SUPPORTED_QUANTIZATION_METHODS}"
-        )
-
-    config_cls = _QUANT_CONFIG_REGISTRY[quantization]
-    logger.info("Creating diffusion quantization config: %s", quantization)
-    return config_cls(**kwargs)
-
-
-def get_vllm_quant_config_for_layers(
-    diffusion_quant_config: DiffusionQuantizationConfig | None,
-) -> "QuantizationConfig | None":
-    """Get the vLLM QuantizationConfig to pass to linear layers.
-
-    This extracts the underlying vLLM config from a DiffusionQuantizationConfig,
-    which can then be passed to vLLM linear layers (QKVParallelLinear, etc.).
-
-    Args:
-        diffusion_quant_config: The diffusion quantization config, or None
-
-    Returns:
-        vLLM QuantizationConfig instance, or None if input is None
-    """
-    if diffusion_quant_config is None:
-        return None
-    return diffusion_quant_config.get_vllm_quant_config()
-
-
-__all__ = [
-    "DiffusionQuantizationConfig",
-    "DiffusionFp8Config",
-    "DiffusionInt8Config",
-    "DiffusionGgufConfig",
-    "get_diffusion_quant_config",
-    "get_vllm_quant_config_for_layers",
-    "SUPPORTED_QUANTIZATION_METHODS",
-]

vllm_omni/diffusion/quantization/base.py

@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Base class for diffusion model quantization configurations."""
-
-from abc import ABC
-from typing import TYPE_CHECKING, ClassVar
-
-import torch
-
-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization.base_config import (
-        QuantizationConfig,
-    )
-
-
-class DiffusionQuantizationConfig(ABC):
-    """Base class for diffusion model quantization configurations.
-
-    This provides a thin wrapper over vLLM's quantization configs,
-    allowing diffusion-model-specific defaults and future extensibility.
-
-    Subclasses should:
-        - Set quant_config_cls to the vLLM QuantizationConfig class
-        - Call super().__init__() after creating self._vllm_config
-        - Optionally override get_name() and get_min_capability() if needed
-    """
-
-    # Subclasses should set this to the vLLM QuantizationConfig class
-    quant_config_cls: ClassVar[type["QuantizationConfig"] | None] = None
-
-    # The underlying vLLM config instance
-    _vllm_config: "QuantizationConfig | None" = None
-
-    @classmethod
-    def get_name(cls) -> str:
-        """Return the quantization method name (e.g., 'fp8', 'int8').
-
-        Delegates to the underlying vLLM config class's get_name().
-        """
-        if cls.quant_config_cls is not None:
-            try:
-                return cls.quant_config_cls.get_name()
-            except TypeError as exc:
-                # Upstream quantization configs are not fully consistent:
-                # some expose get_name as a classmethod (e.g. FP8), while
-                # others expose it as an instance method (e.g. GGUF in
-                # vLLM 0.18.0). Fall back to a default-constructed instance
-                # so diffusion wrappers work with both styles.
-                if "required positional argument" not in str(exc):
-                    raise
-                return cls.quant_config_cls().get_name()
-        raise NotImplementedError("Subclass must set quant_config_cls or override get_name().")
-
-    def get_vllm_quant_config(self) -> "QuantizationConfig | None":
-        """Return the underlying vLLM QuantizationConfig for linear layers."""
-        return self._vllm_config
-
-    @classmethod
-    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
-        """Return supported activation dtypes."""
-        return [torch.bfloat16, torch.float16]
-
-    @classmethod
-    def get_min_capability(cls) -> int:
-        """Minimum GPU compute capability required.
-
-        By default, delegates to the underlying vLLM config class.
-        """
-        if cls.quant_config_cls is not None:
-            return cls.quant_config_cls.get_min_capability()
-        return 80  # Ampere default

vllm_omni/diffusion/quantization/fp8.py

@@ -1,50 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""FP8 quantization config for diffusion transformers."""
-
-from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
-from .base import DiffusionQuantizationConfig
-
-
-class DiffusionFp8Config(DiffusionQuantizationConfig):
-    """FP8 quantization config optimized for diffusion transformers.
-
-    Uses dynamic activation scaling (no calibration dataset needed) and
-    online weight quantization from BF16/FP16 checkpoints.
-
-    Device Compatibility:
-        - Turing (SM 75+): Weight-only FP8 via Marlin kernel
-        - Ada/Hopper (SM 89+): Full W8A8 FP8 with native hardware support
-
-    The kernel selection is automatic based on GPU capability.
-
-    Args:
-        activation_scheme: Activation quantization scheme.
-            - "dynamic": Per-token dynamic scaling (default, no calibration)
-            - "static": Single per-tensor scale (requires calibration)
-        weight_block_size: Block size for block-wise weight quantization.
-            Format: [block_n, block_k]. If None, uses per-tensor scaling.
-        ignored_layers: List of layer name patterns to skip quantization.
-    """
-
-    # Tight coupling with vLLM's Fp8Config - delegates get_name() and get_min_capability()
-    quant_config_cls = Fp8Config
-
-    def __init__(
-        self,
-        activation_scheme: str = "dynamic",
-        weight_block_size: list[int] | None = None,
-        ignored_layers: list[str] | None = None,
-    ):
-        self.activation_scheme = activation_scheme
-        self.weight_block_size = weight_block_size
-        self.ignored_layers = ignored_layers or []
-
-        # Create underlying vLLM FP8 config
-        self._vllm_config = Fp8Config(
-            is_checkpoint_fp8_serialized=False,  # Online quantization from BF16
-            activation_scheme=activation_scheme,
-            weight_block_size=weight_block_size,
-            ignored_layers=ignored_layers,
-        )

vllm_omni/model_executor/models/__init__.py

@@ -1,8 +1,10 @@
-from .bagel.bagel import OmniBagelForConditionalGeneration
-from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration
 from .registry import OmniModelRegistry  # noqa: F401
 
+# Model classes are lazily loaded via OmniModelRegistry.
+# Do NOT eagerly import model classes here — it triggers heavy transitive
+# imports (CUDA, pynvml, bitsandbytes, etc.) that crash in subprocess
+# environments used by vLLM's model inspection.
+
 __all__ = [
-    "Qwen3OmniMoeForConditionalGeneration",
-    "OmniBagelForConditionalGeneration",
+    "OmniModelRegistry",
 ]

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -37,6 +37,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3OmniMoeThinkerMultiModalProcessor,
     Qwen3OmniMoeThinkerProcessingInfo,
 )
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
 
 try:
     import flash_attn
@@ -99,9 +100,16 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config
-        talker_config.text_config.rope_parameters = talker_config.text_config.rope_scaling
-        talker_config.text_config.rope_parameters["rope_theta"] = talker_config.text_config.rope_theta
-        self.quant_config = vllm_config.quant_config
+        rope_params = talker_config.text_config.rope_scaling
+        if rope_params is None:
+            # Newer transformers use rope_parameters instead of rope_scaling
+            rope_params = getattr(talker_config.text_config, "rope_parameters", None) or {}
+        rope_params["rope_theta"] = talker_config.text_config.rope_theta
+        talker_config.text_config.rope_parameters = rope_params
+        quant_config = vllm_config.quant_config
+        if isinstance(quant_config, ComponentQuantizationConfig):
+            quant_config = quant_config.resolve("talker")
+        self.quant_config = quant_config
         self.prefix = prefix
         self.vllm_config = vllm_config
         self.config = talker_config

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -23,6 +23,7 @@
 """Inference-only Qwen3-Omni-Moe model (thinker part)."""
 
 from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import replace
 from functools import partial
 from typing import Any, Literal, cast
 
@@ -47,10 +48,18 @@ from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
 from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.inputs.data import PromptType
 from vllm.logger import init_logger
+from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
 from vllm.model_executor.layers.attention.mm_encoder_attention import (
     MMEncoderAttention,
 )
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.interfaces import (
     MultiModalEmbeddings,
@@ -77,7 +86,15 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import (
     Qwen3Omni_VisionTransformer as _Qwen3Omni_VisionTransformer,
 )
 from vllm.model_executor.models.qwen3_omni_moe_thinker import (
-    Qwen3OmniMoeAudioEncoder,
+    Qwen3OmniMoeAudioAttention as _Qwen3OmniMoeAudioAttention,
+)
+from vllm.model_executor.models.qwen3_omni_moe_thinker import (
+    Qwen3OmniMoeAudioEncoder as _Qwen3OmniMoeAudioEncoder,
+)
+from vllm.model_executor.models.qwen3_omni_moe_thinker import (
+    Qwen3OmniMoeAudioEncoderLayer as _Qwen3OmniMoeAudioEncoderLayer,
+)
+from vllm.model_executor.models.qwen3_omni_moe_thinker import (
     _get_feat_extract_output_lengths,
 )
 from vllm.model_executor.models.utils import (
@@ -102,6 +119,7 @@ from vllm.transformers_utils.processor import cached_processor_from_config
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
 )
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
 
 try:
     import flash_attn
@@ -216,6 +234,252 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer):
         return hidden_states
 
 
+class Qwen3OmniMoeAudioAttention(_Qwen3OmniMoeAudioAttention):
+    """Subclass that adds quant_config support to audio attention."""
+
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        nn.Module.__init__(self)
+        self.embed_dim = config.d_model
+        self.num_heads = config.encoder_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_local_heads = self.num_heads // tp_size
+
+        if (self.head_dim * self.num_heads) != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: "
+                f"{self.embed_dim} and `num_heads`: {self.num_heads})."
+            )
+
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv = QKVParallelLinear(
+            hidden_size=self.embed_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            total_num_kv_heads=self.num_heads,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv",
+        )
+
+        self.out_proj = RowParallelLinear(
+            input_size=self.embed_dim,
+            output_size=self.embed_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+        )
+
+        self.attn = MMEncoderAttention(
+            num_heads=self.num_local_heads,
+            head_size=self.head_dim,
+            scale=self.scaling,
+            prefix=f"{prefix}.attn",
+        )
+
+
+class Qwen3OmniMoeAudioEncoderLayer(_Qwen3OmniMoeAudioEncoderLayer):
+    """Subclass that adds quant_config support to audio encoder layers."""
+
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        nn.Module.__init__(self)
+        self.embed_dim = config.d_model
+        self.self_attn = Qwen3OmniMoeAudioAttention(config, quant_config=quant_config, prefix=f"{prefix}.self_attn")
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function]
+        self.fc1 = ColumnParallelLinear(
+            self.embed_dim,
+            config.encoder_ffn_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc1",
+        )
+        self.fc2 = RowParallelLinear(
+            config.encoder_ffn_dim,
+            self.embed_dim,
+            bias=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc2",
+        )
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+
+class Qwen3OmniMoeAudioEncoder(_Qwen3OmniMoeAudioEncoder):
+    """Subclass that adds quant_config support to the audio encoder."""
+
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        nn.Module.__init__(self)
+
+        embed_dim = config.d_model
+        self.num_mel_bins = config.num_mel_bins
+        self.max_source_positions = config.max_source_positions
+        self.n_window = config.n_window
+        self.n_window_infer = config.n_window_infer
+        self.conv_chunksize = config.conv_chunksize
+
+        from vllm.model_executor.models.qwen3_omni_moe_thinker import (
+            SinusoidsPositionEmbedding,
+        )
+
+        self.positional_embedding = SinusoidsPositionEmbedding(self.max_source_positions, embed_dim)
+
+        # Convolutional layers (not quantized)
+        self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)
+        self.conv2d2 = nn.Conv2d(
+            config.downsample_hidden_size,
+            config.downsample_hidden_size,
+            3,
+            2,
+            padding=1,
+        )
+        self.conv2d3 = nn.Conv2d(
+            config.downsample_hidden_size,
+            config.downsample_hidden_size,
+            3,
+            2,
+            padding=1,
+        )
+
+        conv_out_dim = config.downsample_hidden_size * ((((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2)
+        self.conv_out = ReplicatedLinear(
+            conv_out_dim,
+            config.d_model,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.conv_out",
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                Qwen3OmniMoeAudioEncoderLayer(
+                    config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{i}",
+                )
+                for i in range(config.encoder_layers)
+            ]
+        )
+
+        self.ln_post = nn.LayerNorm(config.d_model)
+        self.proj1 = ReplicatedLinear(
+            config.d_model,
+            config.d_model,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj1",
+        )
+        self.act = _ACTIVATION_REGISTRY[config.activation_function]
+        self.proj2 = ReplicatedLinear(
+            config.d_model,
+            config.output_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj2",
+        )
+
+        from vllm.model_executor.models.vision import get_vit_attn_backend
+
+        self.attn_backend = get_vit_attn_backend(
+            head_size=config.d_model // config.encoder_attention_heads,
+            dtype=torch.get_default_dtype(),
+        )
+
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        feature_lens: torch.Tensor,
+        aftercnn_lens: torch.Tensor,
+    ):
+        import torch.nn.functional as F
+
+        chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()
+
+        chunk_lengths = torch.tensor(
+            [self.n_window * 2] * chunk_num.sum(),
+            dtype=torch.long,
+            device=feature_lens.device,
+        )
+        tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
+        chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)
+        chunk_lengths[chunk_lengths == 0] = self.n_window * 2
+
+        chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
+        padded_feature = nn.utils.rnn.pad_sequence(chunk_list, batch_first=True).transpose(1, 2)
+
+        feature_lens_after_cnn = self._get_cnn_output_lengths(chunk_lengths)
+        max_len_after_cnn = feature_lens_after_cnn.max().item()
+        indices = torch.arange(max_len_after_cnn, device=padded_feature.device)
+        padded_mask_after_cnn = indices.unsqueeze(0) < feature_lens_after_cnn.unsqueeze(1)
+
+        padded_feature = padded_feature.unsqueeze(1)
+
+        if padded_feature.size(0) <= self.conv_chunksize:
+            padded_embed = F.gelu(self.conv2d1(padded_feature))
+            padded_embed = F.gelu(self.conv2d2(padded_embed))
+            padded_embed = F.gelu(self.conv2d3(padded_embed))
+        else:
+            padded_embeds = []
+            for chunk in padded_feature.split(self.conv_chunksize, dim=0):
+                padded_embed = F.gelu(self.conv2d1(chunk))
+                padded_embed = F.gelu(self.conv2d2(padded_embed))
+                padded_embed = F.gelu(self.conv2d3(padded_embed))
+                padded_embeds.append(padded_embed)
+            padded_embed = torch.cat(padded_embeds, dim=0)
+
+        b, c, f, t = padded_embed.size()
+        padded_embed, _ = self.conv_out(padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f))
+
+        positional_embedding = (
+            self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]
+            .unsqueeze(0)
+            .to(padded_embed.dtype)
+        )
+        padded_embed = padded_embed + positional_embedding
+
+        hidden_states = padded_embed[padded_mask_after_cnn]
+
+        cu_chunk_lens = [0]
+        window_aftercnn = padded_mask_after_cnn.shape[-1] * (self.n_window_infer // (self.n_window * 2))
+        for cnn_len in aftercnn_lens.tolist():
+            num_full_chunks = cnn_len // window_aftercnn
+            remainder = cnn_len % window_aftercnn
+            cu_chunk_lens.extend([window_aftercnn] * num_full_chunks)
+            if remainder:
+                cu_chunk_lens.append(remainder)
+        cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(-1, dtype=torch.int32)
+
+        max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
+
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer(
+                hidden_states,
+                cu_seqlens,
+                max_seqlen,
+            )
+
+        # ReplicatedLinear returns (output, bias) — unpack
+        hidden_states = self.ln_post(hidden_states)
+        hidden_states, _ = self.proj1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states, _ = self.proj2(hidden_states)
+
+        return hidden_states
+
+
 # Speech input languages supported by Qwen3-Omni
 # From: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct
 ISO639_1_SUPPORTED_LANGS = {
@@ -850,9 +1114,41 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
         self.multimodal_config = multimodal_config
         self.quant_config = quant_config
 
+        # Pre-quantized checkpoints (modelopt NVFP4/FP8/MXFP8) quantize the
+        # entire thinker — audio tower, visual encoder, and language model
+        # all share the same quant method.  Dynamic quantization methods
+        # (e.g. --quantization fp8) should only target the language model.
+        _PRE_QUANTIZED_METHODS = {"modelopt", "modelopt_fp4", "modelopt_mxfp8"}
+
+        if isinstance(quant_config, ComponentQuantizationConfig):
+            audio_quant_config = quant_config.resolve("audio_tower")
+            visual_quant_config = quant_config.resolve("visual")
+            language_quant_config = quant_config.resolve("language_model")
+        elif quant_config is not None:
+            if quant_config.get_name() in _PRE_QUANTIZED_METHODS:
+                # Pre-quantized: pass quant_config to all subcomponents.
+                audio_quant_config = quant_config
+                visual_quant_config = quant_config
+                language_quant_config = quant_config
+            else:
+                # Dynamic quantization: scope to language_model only.
+                quant_config = ComponentQuantizationConfig(
+                    component_configs={"language_model": quant_config},
+                    default_config=None,
+                )
+                vllm_config = replace(vllm_config, quant_config=quant_config)
+                audio_quant_config = None
+                visual_quant_config = None
+                language_quant_config = quant_config.resolve("language_model")
+        else:
+            audio_quant_config = None
+            visual_quant_config = None
+            language_quant_config = None
+
         with self._mark_tower_model(vllm_config, "audio"):
             self.audio_tower = Qwen3OmniMoeAudioEncoder(
                 thinker_config.audio_config,
+                quant_config=audio_quant_config,
                 prefix=maybe_prefix(prefix, "audio_tower"),
             )
 
@@ -867,7 +1163,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
             self.visual = Qwen3Omni_VisionTransformer(
                 vision_config=thinker_config.vision_config,
                 norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
-                quant_config=quant_config,
+                quant_config=visual_quant_config,
                 prefix=maybe_prefix(prefix, "visual"),
             )
 
@@ -882,11 +1178,14 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
                 ]
 
         with self._mark_language_model(vllm_config):
+            lm_vllm_config = vllm_config.with_hf_config(
+                thinker_config.text_config,
+                architectures=["Qwen3MoeForCausalLM"],
+            )
+            if language_quant_config is not quant_config:
+                lm_vllm_config = replace(lm_vllm_config, quant_config=language_quant_config)
             self.language_model = Qwen3MoeLLMForCausalLM(
-                vllm_config=vllm_config.with_hf_config(
-                    thinker_config.text_config,
-                    architectures=["Qwen3MoeForCausalLM"],
-                ),
+                vllm_config=lm_vllm_config,
                 prefix=maybe_prefix(prefix, "language_model"),
             )
 

vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit_2gpu_fp8.yaml

@@ -0,0 +1,39 @@
+# Stage config for running Hunyuan-Image3.0 DiT with FP8 online quantization.
+# The following config is for 2x H200 GPU.
+
+# Stage 0:  Diffusion (DiT + VAE)
+# This stage receives noise and timesteps and performs denoising + VAE decode
+stage_args:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      quantization: "fp8"
+      parallel_config:
+        tensor_parallel_size: 2
+        enable_expert_parallel: true
+      omni_kv_config:
+        need_recv_cache: true
+
+    final_output: true
+    final_output_type: image
+    is_comprehension: false
+    default_sampling_params:
+      seed: 42
+
+# Runtime edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1

vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe_2gpu.yaml

@@ -0,0 +1,44 @@
+# Stage config for running Hunyuan-Image3.0 on 2 GPUs with FP8.
+# Stage 0: AR Model (vLLM implementation)
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0,1"
+    engine_args:
+      model_stage: AR
+      max_num_seqs: 1
+      model_arch: HunyuanImage3ForCausalMM
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 2
+      pipeline_parallel_size: 1
+      hf_overrides:
+        rope_parameters:
+          mrope_section: [0, 32, 32]
+          rope_type: default
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1

vllm_omni/quantization/__init__.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unified quantization framework for vLLM-OMNI.
+
+Delegates to vLLM's quantization registry (35+ methods, all platforms).
+Adds per-component quantization for multi-stage models.
+
+    from vllm_omni.quantization import build_quant_config
+
+    config = build_quant_config("fp8")
+    config = build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
+"""
+
+from .component_config import ComponentQuantizationConfig
+from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config
+
+# DiffusionGGUFConfig is NOT imported here to avoid pulling in
+# GGUF -> fused_moe -> pynvml at module load time.
+
+__all__ = [
+    "build_quant_config",
+    "ComponentQuantizationConfig",
+    "SUPPORTED_QUANTIZATION_METHODS",
+]

vllm_omni/quantization/component_config.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-component quantization routing for multi-stage models.
+
+Routes get_quant_method() to different configs based on longest-prefix match:
+    {"transformer": fp8_config, "vae": None}
+    "transformer.blocks.0.attn.to_q" -> fp8_config
+    "vae.encoder.conv_in"            -> None
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+)
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizeMethodBase,
+    )
+
+
+class ComponentQuantizationConfig(QuantizationConfig):
+    """Routes quantization to different configs by layer prefix."""
+
+    def __init__(
+        self,
+        component_configs: dict[str, QuantizationConfig | None],
+        default_config: QuantizationConfig | None = None,
+    ) -> None:
+        self._components = component_configs
+        self._default = default_config
+        self._sorted_prefixes = sorted(self._components.keys(), key=len, reverse=True)
+
+    def resolve(self, prefix: str) -> QuantizationConfig | None:
+        """Find the config for a given layer prefix (longest-prefix match).
+
+        Note: vLLM may remap quantization prefixes vs model definition
+        prefixes (e.g. via WeightsMapper). If prefixes don't match after
+        remapping, layers may fall through to the default config.
+        """
+        for comp_prefix in self._sorted_prefixes:
+            if prefix.startswith(comp_prefix):
+                return self._components[comp_prefix]
+        return self._default
+
+    def get_name(self) -> str:
+        return "component"
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
+        config = self.resolve(prefix)
+        if config is None:
+            return None
+        return config.get_quant_method(layer, prefix)
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    def get_min_capability(self) -> int:
+        """Return the minimum capability across all component configs."""
+        caps = [c.get_min_capability() for c in self._components.values() if c is not None]
+        if self._default is not None:
+            caps.append(self._default.get_min_capability())
+        return min(caps) if caps else 0
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> ComponentQuantizationConfig:
+        raise NotImplementedError("Use build_quant_config() instead")
+
+    def get_config_filenames(self) -> list[str]:
+        return []
+
+    @property
+    def component_configs(self) -> dict[str, QuantizationConfig | None]:
+        return self._components
+
+    @property
+    def default_config(self) -> QuantizationConfig | None:
+        return self._default

vllm_omni/quantization/factory.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Factory for building quantization configs.
+
+build_quant_config() delegates to vLLM's quantization registry.
+The only extension point is _OVERRIDES for methods that need a
+different QuantizationConfig subclass in the OMNI context (e.g. GGUF).
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable, Mapping
+from typing import Any
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization import (
+    QUANTIZATION_METHODS,
+    get_quantization_config,
+)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+)
+
+from .component_config import ComponentQuantizationConfig
+
+logger = init_logger(__name__)
+
+
+def _build_gguf(**kw: Any) -> QuantizationConfig:
+    """Lazy import to avoid pulling in CUDA/pynvml at module load time."""
+    from .gguf_config import DiffusionGGUFConfig
+
+    return DiffusionGGUFConfig(**kw)
+
+
+def _build_int8(**kw: Any) -> QuantizationConfig:
+    """Lazy import for Int8 diffusion config (supports CUDA + NPU)."""
+    from .int8_config import DiffusionInt8Config
+
+    return DiffusionInt8Config(**kw)
+
+
+_OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
+    "gguf": _build_gguf,
+    "int8": _build_int8,
+}
+
+SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
+
+
+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()
+
+    if method in _OVERRIDES:
+        return _OVERRIDES[method](**kwargs)
+
+    if method not in QUANTIZATION_METHODS:
+        raise ValueError(f"Unknown quantization method: {method!r}. Supported: {SUPPORTED_QUANTIZATION_METHODS}")
+
+    config_cls = get_quantization_config(method)
+
+    try:
+        return config_cls(**kwargs)
+    except TypeError:
+        sig = inspect.signature(config_cls.__init__)
+        raise TypeError(
+            f"Cannot instantiate {config_cls.__name__} with kwargs {kwargs}. Expected signature: {sig}"
+        ) from None
+
+
+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
+    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".
+    """
+    if "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())
+
+
+def _build_component_config(spec: dict[str, Any]) -> ComponentQuantizationConfig:
+    """Build ComponentQuantizationConfig from a per-component dict."""
+    component_configs: dict[str, QuantizationConfig | None] = {}
+    default_config: QuantizationConfig | None = None
+
+    for prefix, value in spec.items():
+        if value is None:
+            config = None
+        elif isinstance(value, str):
+            config = _build_single(value)
+        elif isinstance(value, dict):
+            value = dict(value)  # avoid mutating caller's dict
+            method = value.pop("method", None)
+            if method is None:
+                raise ValueError(f"Component '{prefix}' config dict must have a 'method' key")
+            config = _build_single(method, **value)
+        else:
+            raise TypeError(f"Component '{prefix}' config must be str, dict, or None, got {type(value).__name__}")
+
+        if prefix == "default":
+            default_config = config
+        else:
+            component_configs[prefix] = config
+
+    logger.info(
+        "Per-component quantization: %s",
+        {k: (v.get_name() if v else None) for k, v in component_configs.items()},
+    )
+    return ComponentQuantizationConfig(component_configs, default_config)
+
+
+def build_quant_config(
+    spec: str | dict[str, Any] | QuantizationConfig | None,
+    **kwargs: Any,
+) -> QuantizationConfig | None:
+    """Build a quantization config from a flexible specification.
+
+    Args:
+        spec: None/"none", method name str, dict with "method" key,
+              per-component dict, or QuantizationConfig passthrough.
+        **kwargs: Extra params merged with dict spec.
+    """
+    if spec is None:
+        return None
+
+    if isinstance(spec, QuantizationConfig):
+        return spec
+
+    if isinstance(spec, str):
+        if spec.lower() == "none":
+            return None
+        logger.info("Building quantization config: %s", spec)
+        return _build_single(spec, **kwargs)
+
+    if isinstance(spec, Mapping):
+        spec = dict(spec)
+
+        if _is_per_component_dict(spec):
+            return _build_component_config(spec)
+
+        method = spec.pop("method", None)
+        if method is None:
+            raise ValueError(
+                "Dict quantization config must have a '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)
+
+    raise TypeError(f"quantization config must be str, dict, QuantizationConfig, or None, got {type(spec).__name__}")

vllm_omni/quantization/gguf_config.py

@@ -1,6 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""GGUF quantization config for diffusion transformers."""
+"""GGUF quantization config for diffusion transformers.
+
+Uses dequant+GEMM instead of the fused kernel path (which expects 2D inputs).
+"""
+
+from __future__ import annotations
 
 import gguf
 import torch
@@ -15,8 +20,6 @@ from vllm.model_executor.layers.quantization.gguf import (
     is_layer_skipped_gguf,
 )
 
-from .base import DiffusionQuantizationConfig
-
 
 def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor:
     if qweight_type in UNQUANTIZED_TYPES:
@@ -28,19 +31,17 @@ def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int)
 
 
 class DiffusionGGUFLinearMethod(GGUFLinearMethod):
+    """GGUF linear method using dequant+GEMM for N-D diffusion tensors."""
+
     def apply(
         self,
         layer: torch.nn.Module,
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        # Dequantize + GEMM path: torch.matmul multiplies over the last
-        # dimension and broadcasts leading dimensions, so no 2D flattening
-        # is required here.
         shard_id = getattr(layer.qweight, "shard_id", None)
 
         if shard_id:
-            # dequantize shard weights respectively
             shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id
             qweight = layer.qweight
             result = []
@@ -53,41 +54,26 @@ class DiffusionGGUFLinearMethod(GGUFLinearMethod):
             qweight = layer.qweight
             qweight_type = layer.qweight_type.weight_type
             out = dequant_gemm_gguf(x, qweight, qweight_type)
+
         if bias is not None:
             out.add_(bias)
         return out
 
 
-class _GGUFConfig(GGUFConfig):
-    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> "QuantizeMethodBase":
-        if isinstance(layer, LinearBase):
-            if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
-                return UnquantizedLinearMethod()
-            return DiffusionGGUFLinearMethod(self)
-        return None
-
-
-class DiffusionGgufConfig(DiffusionQuantizationConfig):
-    """GGUF quantization config for diffusion transformers.
-
-    This is a thin wrapper around vLLM's GGUFConfig and also carries
-    the GGUF model reference for loader use.
-
-    Args:
-        gguf_model: GGUF model path or HF reference (repo/file or repo:quant_type)
-        unquantized_modules: Optional list of module name patterns to skip GGUF
-            quantization. Note: diffusion linear layers often use short prefixes
-            (e.g., "to_qkv"), so these patterns are matched as substrings.
-    """
-
-    quant_config_cls = GGUFConfig
+class DiffusionGGUFConfig(GGUFConfig):
+    """GGUF config that carries gguf_model path and uses dequant+GEMM."""
 
     def __init__(
         self,
         gguf_model: str | None = None,
         unquantized_modules: list[str] | None = None,
     ) -> None:
+        super().__init__(unquantized_modules=unquantized_modules or [])
         self.gguf_model = gguf_model
-        self.unquantized_modules = unquantized_modules or []
 
-        self._vllm_config = _GGUFConfig(unquantized_modules=self.unquantized_modules)
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
+                return UnquantizedLinearMethod()
+            return DiffusionGGUFLinearMethod(self)
+        return None

vllm_omni/quantization/int8_config.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""INT8 quantization config for diffusion transformers."""
+"""INT8 quantization config for diffusion transformers.
+
+Supports both online (dynamic) and offline (checkpoint) INT8 quantization
+on CUDA and NPU platforms.
+"""
 
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, Optional
@@ -40,8 +44,6 @@ if current_omni_platform.is_npu():
 else:
     torch_npu = None
 
-from .base import DiffusionQuantizationConfig
-
 if TYPE_CHECKING:
     from vllm.model_executor.models.utils import WeightsMapper
 
@@ -72,9 +74,12 @@ def create_weight_parameter(
     )
 
 
-class Int8Config(QuantizationConfig):
-    """
-    Config class for Int8.
+class DiffusionInt8Config(QuantizationConfig):
+    """INT8 quantization config for diffusion transformers.
+
+    Supports online (dynamic) quantization from BF16/FP16 checkpoints
+    and offline quantization from serialized INT8 checkpoints.
+    Works on both CUDA and NPU platforms.
     """
 
     def __init__(
@@ -114,7 +119,7 @@ class Int8Config(QuantizationConfig):
             self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers)
 
     @classmethod
-    def from_config(cls, config: dict[str, Any]) -> "Int8Config":
+    def from_config(cls, config: dict[str, Any]) -> "DiffusionInt8Config":
         quant_method = cls.get_from_keys(config, ["quant_method"])
         is_checkpoint_int8_serialized = "int8" in quant_method
         activation_scheme = cls.get_from_keys_or(config, ["activation_scheme"], "dynamic")
@@ -168,7 +173,7 @@ class BaseInt8LinearMethod(LinearMethodBase):
         quant_config: The quantization config.
     """
 
-    def __init__(self, quant_config: Int8Config):
+    def __init__(self, quant_config: DiffusionInt8Config):
         self.quant_config = quant_config
         self.out_dtype = torch.get_default_dtype()
 
@@ -315,7 +320,7 @@ class Int8LinearMethod(BaseInt8LinearMethod):
         quant_config: The quantization config.
     """
 
-    def __init__(self, quant_config: Int8Config):
+    def __init__(self, quant_config: DiffusionInt8Config):
         super().__init__(quant_config)
 
         self.int8_linear = init_int8_linear_kernel(
@@ -346,7 +351,7 @@ class NPUInt8LinearMethod(BaseInt8LinearMethod):
         quant_config: The quantization config.
     """
 
-    def __init__(self, quant_config: Int8Config):
+    def __init__(self, quant_config: DiffusionInt8Config):
         super().__init__(quant_config)
 
     def process_weights_after_loading(self, layer: Module) -> None:
@@ -445,31 +450,3 @@ class NPUInt8OnlineLinearMethod(LazyWeightMixin, NPUInt8LinearMethod):
 
         # Prevent duplicate processing (e.g., during weight reload)
         layer._already_called_process_weights_after_loading = True
-
-
-class DiffusionInt8Config(DiffusionQuantizationConfig):
-    """
-    Int8 quantization config optimized for diffusion transformers.
-
-    Args:
-        activation_scheme: Activation quantization scheme.
-            - "dynamic": Per-token dynamic scaling (default, no calibration)
-        ignored_layers: List of layer name patterns to skip quantization.
-    """
-
-    quant_config_cls = Int8Config
-
-    def __init__(
-        self,
-        activation_scheme: str = "dynamic",
-        ignored_layers: list[str] | None = None,
-    ):
-        self.activation_scheme = activation_scheme
-        self.ignored_layers = ignored_layers or []
-
-        # Create underlying vLLM Int8 config
-        self._vllm_config = Int8Config(
-            is_checkpoint_int8_serialized=False,  # Online quantization
-            activation_scheme=activation_scheme,
-            ignored_layers=ignored_layers,
-        )