diff --git a/benchmarks/diffusion/quantization_quality.py b/benchmarks/diffusion/quantization_quality.py new file mode 100644 index 000000000..4a916e7ea --- /dev/null +++ b/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 + / # 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) diff --git a/docs/.nav.yml b/docs/.nav.yml index 1d7346669..48917d526 100644 --- a/docs/.nav.yml +++ b/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 diff --git a/docs/contributing/model/adding_quantization_model.md b/docs/contributing/model/adding_quantization_model.md new file mode 100644 index 000000000..f27318888 --- /dev/null +++ b/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 diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py index 5acd1891e..3214abd56 100644 --- a/docs/mkdocs/hooks/generate_api_readme.py +++ b/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", diff --git a/docs/user_guide/diffusion/quantization/fp8.md b/docs/user_guide/diffusion/quantization/fp8.md index f7887ab76..338781fa2 100644 --- a/docs/user_guide/diffusion/quantization/fp8.md +++ b/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 diff --git a/docs/user_guide/diffusion/quantization/gguf.md b/docs/user_guide/diffusion/quantization/gguf.md index 04e2d6aad..f3bb17b85 100644 --- a/docs/user_guide/diffusion/quantization/gguf.md +++ b/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. diff --git a/docs/user_guide/diffusion/quantization/overview.md b/docs/user_guide/diffusion/quantization/overview.md index a95afdbf4..0fc8b9bc2 100644 --- a/docs/user_guide/diffusion/quantization/overview.md +++ b/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, +}) +``` diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index ae384cf37..6562f32ae 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/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) diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 793908519..17faf5f15 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/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", diff --git a/mkdocs.yml b/mkdocs.yml index c81d384d9..cab033920 100644 --- a/mkdocs.yml +++ b/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 diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py index 57661d299..9c18c1f55 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/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" diff --git a/tests/diffusion/quantization/test_gguf_config.py b/tests/diffusion/quantization/test_gguf_config.py index bd3d1fcb2..1dd401d59 100644 --- a/tests/diffusion/quantization/test_gguf_config.py +++ b/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( [ diff --git a/tests/diffusion/quantization/test_int8_config.py b/tests/diffusion/quantization/test_int8_config.py index 9b5d67fcb..d4d5aa5a7 100644 --- a/tests/diffusion/quantization/test_int8_config.py +++ b/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) diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py new file mode 100644 index 000000000..a937a6486 --- /dev/null +++ b/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}" + ) diff --git a/tests/e2e/offline_inference/compute_lpips.py b/tests/e2e/offline_inference/compute_lpips.py new file mode 100644 index 000000000..4d87a6dc8 --- /dev/null +++ b/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: _bf16.png and _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() diff --git a/tests/e2e/offline_inference/run_quantization_e2e.sh b/tests/e2e/offline_inference/run_quantization_e2e.sh new file mode 100644 index 000000000..8bd921cfb --- /dev/null +++ b/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 diff --git a/tests/e2e/offline_inference/test_quantization_fp8.py b/tests/e2e/offline_inference/test_quantization_fp8.py new file mode 100644 index 000000000..5943afa02 --- /dev/null +++ b/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 diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index db558f64a..588efabfc 100644 --- a/vllm_omni/config/model.py +++ b/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, diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 98d8e8c83..488378b40 100644 --- a/vllm_omni/diffusion/data.py +++ b/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": } 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: diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index e2c6e2341..c48640e34 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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": diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 367464692..685d14729 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 29ac411f5..c4155a9fc 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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 ) diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index 3dde14866..6f43e8dbb 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/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( diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 6261ff000..437dd58d0 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/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) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py index 2fe486714..4dbe44fbb 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py +++ b/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 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 49e6f780b..505bad3d5 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/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) diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 9d2ed036e..ac18d5773 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/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() diff --git a/vllm_omni/diffusion/quantization/__init__.py b/vllm_omni/diffusion/quantization/__init__.py deleted file mode 100644 index eb4b8ea9f..000000000 --- a/vllm_omni/diffusion/quantization/__init__.py +++ /dev/null @@ -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", -] diff --git a/vllm_omni/diffusion/quantization/base.py b/vllm_omni/diffusion/quantization/base.py deleted file mode 100644 index f008eca0b..000000000 --- a/vllm_omni/diffusion/quantization/base.py +++ /dev/null @@ -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 diff --git a/vllm_omni/diffusion/quantization/fp8.py b/vllm_omni/diffusion/quantization/fp8.py deleted file mode 100644 index 68abf9c22..000000000 --- a/vllm_omni/diffusion/quantization/fp8.py +++ /dev/null @@ -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, - ) diff --git a/vllm_omni/model_executor/models/__init__.py b/vllm_omni/model_executor/models/__init__.py index 68074bbb9..065480078 100644 --- a/vllm_omni/model_executor/models/__init__.py +++ b/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", ] diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 7b3d0b017..bb491d01b 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/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 diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 4693a2314..fd7b14ab4 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/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"), ) diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit_2gpu_fp8.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit_2gpu_fp8.yaml new file mode 100644 index 000000000..51110c285 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe_2gpu.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe_2gpu.yaml new file mode 100644 index 000000000..e029c3836 --- /dev/null +++ b/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 diff --git a/vllm_omni/quantization/__init__.py b/vllm_omni/quantization/__init__.py new file mode 100644 index 000000000..b50fe8607 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py new file mode 100644 index 000000000..7986da885 --- /dev/null +++ b/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 diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py new file mode 100644 index 000000000..a867f37a4 --- /dev/null +++ b/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__}") diff --git a/vllm_omni/diffusion/quantization/gguf.py b/vllm_omni/quantization/gguf_config.py similarity index 66% rename from vllm_omni/diffusion/quantization/gguf.py rename to vllm_omni/quantization/gguf_config.py index 85ec2ee33..42b2b2fe2 100644 --- a/vllm_omni/diffusion/quantization/gguf.py +++ b/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 diff --git a/vllm_omni/diffusion/quantization/int8.py b/vllm_omni/quantization/int8_config.py similarity index 92% rename from vllm_omni/diffusion/quantization/int8.py rename to vllm_omni/quantization/int8_config.py index 798038c9f..37d430047 100644 --- a/vllm_omni/diffusion/quantization/int8.py +++ b/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, - )