# 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] if hasattr(first, "images") and first.images: return first.images[0], peak_mem inner = first.request_output if inner is not None and hasattr(inner, "images") and inner.images: return inner.images[0], peak_mem raise ValueError("Could not extract image from output.") 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}" )