diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 503f369d8..19af2f484 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -139,7 +139,7 @@ The following tables show which models support each feature: | **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **Qwen-Image-2.1** | ❌ | ❌ | ✅ (Ulysses only) | ✅ | ❓ | ❌ | ❓ | ❓ | ✅ (decode) | ❓ | ✅ | +| **Qwen-Image-2.1** | ❌ | ❌ | ✅ (Ulysses only) | ✅ | ❓ | ❌ | ❓ | ❓ | ✅ (decode) | ✅ | ✅ | | **SenseNova-U1 / U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 010ef4ba7..71c63ac95 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -1,4 +1,4 @@ -# FP8 Quantization +# FP8 Quantization ## Overview @@ -85,6 +85,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) | Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization | |-------|-----------|:-------:|:------:|----------------|------------------|------------------| | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | | +| Qwen-Image-2.1 | `Qwen/Qwen-Image-2.1` | Yes | Not validated | Skip image-stream MLPs for best fidelity; all-layer FP8 is loadable but composition drifts on some prompts. On Blackwell GB200 FP8 saves ~17% peak memory but is not faster than BF16 at 1024x1024. Text encoder FP8 covers the language model only — the vision tower and LM head stay BF16 | `img_mlp` | ✅︎ (language model only) | | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | | | LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | | LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | | diff --git a/recipes/Qwen/Qwen-Image-2.1.md b/recipes/Qwen/Qwen-Image-2.1.md index 038a7158b..f833ac58e 100644 --- a/recipes/Qwen/Qwen-Image-2.1.md +++ b/recipes/Qwen/Qwen-Image-2.1.md @@ -183,7 +183,81 @@ automatically when the checkpoint declares `causal_condition: true` and is re-created per generation, so no state leaks between requests. It changes the compute path — cached prefix vs. full recompute at every step — so exact numerical reproduction against a no-cache reference (e.g. the diffusers -pipeline) requires matching the cache setting on both sides. +pipeline) requires matching the cache setting on both sides. The prefix KV +cache is orthogonal to weight quantization and works unchanged under FP8. + +### FP8 quantization + +Online FP8 quantization of the transformer is supported; the fused +`to_qkv` projections load correctly, and the prefix KV cache path is +unaffected. Only block-internal linears are eligible — the boundary +projections (`img_in`/`txt_in`/`modulation`/`proj_out`/`time_text_embed`) stay +in BF16 by construction: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model Qwen/Qwen-Image-2.1 \ + --prompt "A ceramic teapot on a wooden table" \ + --negative-prompt "blurry, low quality, text, watermark" \ + --output qwen_image_21_fp8.png \ + --num-inference-steps 50 --cfg-scale 4.0 \ + --quantization fp8 --ignored-layers "img_mlp" +``` + +`ignored_layers` entries are matched as name patterns for this model (e.g. +`img_mlp` matches every block's `img_mlp.proj`/`img_mlp.gate_layer`/ +`img_mlp.out`); exact full prefixes like +`transformer.transformer_blocks.3.attn.to_qkv` also work. + +A per-module sensitivity sweep (fixed prompt embeds / latent / timestep, 5 +timesteps across the trajectory, relative L2 of the noise prediction vs BF16) +ranks the groups: `img_mlp` is the most sensitive (3.8% max error on its own, +driven by `img_mlp.proj` 2.8% and `img_mlp.out` 2.5%), then `attn.to_out` +(2.7%), then `attn.to_qkv` (1.6%); errors are largest at the late (low-sigma) +steps. End-to-end on 4 fixed prompts at 1024x1024, 50 steps, seed 42 +(GB200, single GPU): + +| Config | Quantized block linears | Avg PSNR vs BF16 | Min PSNR | Steady-state time/image | Peak memory | +| --- | --- | --- | --- | --- | --- | +| BF16 | 0 / 160 | — | — | 7.2 s | 40.0 GB | +| FP8 all layers | 160 / 160 | 26.1 dB | 19.1 dB | 8.4 s | 33.4 GB | +| FP8, `ignored_layers=["img_mlp"]` | 64 / 160 | 29.7 dB | 21.6 dB | 8.8 s | 38.3 GB | + +All-layer FP8 is usable but its composition can drift on some prompts +(fine detail and layout shift, worst case ~19 dB); keeping `img_mlp` in BF16 +preserves composition noticeably better and is the recommended setting. On +Blackwell (GB200) FP8 here is a memory optimization, not a speedup: dynamic +per-token activation quantization costs more than the FP8 GEMM saves at these +shapes, so BF16 remains the fastest option. + +## Quantization + +Online FP8 is supported per component via `quantization_config`: + +```python +from vllm_omni import Omni + +omni = Omni( + model="Qwen/Qwen-Image-2.1", + quantization_config={ + # DiT: skip the sensitive image-stream MLPs if quality regresses. + "transformer": {"method": "fp8", "ignored_layers": ["img_mlp"]}, + # Qwen3-VL text encoder: only the language-model linear layers are + # quantized; the vision tower and the (unused) LM head stay BF16. + "text_encoder": {"method": "fp8"}, + }, +) +``` + +Either component can be quantized on its own. Measured on GB200 at 1024×1024 +(seed 42, 50 steps): text-encoder FP8 lowers peak GPU memory from ~41.0 GiB to +~34.4 GiB (−6.6 GiB); adding DiT FP8 reaches ~28.1 GiB. Output quality stays +close to BF16 (T2I PSNR vs BF16 ≈ 28–32 dB with text-encoder FP8; the +image-conditioned edit path ≈ 36 dB). FP8 saves memory but does not speed up +generation on this hardware. The edit path routes condition images through the +BF16 vision tower, so it is unaffected by text-encoder FP8. See +[`docs/user_guide/quantization/fp8.md`](../../docs/user_guide/quantization/fp8.md) +for the scope rules. ## Known Limitations diff --git a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py index 6688e0614..ca5bc594f 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py +++ b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py @@ -3,6 +3,7 @@ # Adapted from diffusers' pipeline_qwenimage21.py (Qwen-Image 2.1 T2I + image-conditioned generation). import copy +import dataclasses import json import logging import os @@ -19,7 +20,7 @@ from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( ) from diffusers.utils.torch_utils import randn_tensor from torch import nn -from transformers import Qwen3VLForConditionalGeneration, Qwen3VLProcessor +from transformers import AutoConfig, AutoModelForImageTextToText, Qwen3VLForConditionalGeneration, Qwen3VLProcessor from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -44,6 +45,7 @@ from vllm_omni.diffusion.models.qwen_image_21.cfg_parallel import QwenImage21CFG from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import ( QwenImage21Transformer2DModel, ) +from vllm_omni.diffusion.models.utils import create_transformers_model from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.prompt_utils import ( @@ -55,8 +57,12 @@ from vllm_omni.diffusion.utils.size_utils import ( from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt +from vllm_omni.quantization import resolve_component_quant_config +from vllm_omni.quantization.component_config import resolve_encoder_quant_config if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import StepRequestState @@ -96,6 +102,50 @@ def _read_vae_scale_factor(od_config: OmniDiffusionConfig) -> int: return 16 +def _resolve_text_encoder_quant_config( + quant_config: "QuantizationConfig | None", +) -> "QuantizationConfig | None": + """Resolve the text-encoder slice of a global or per-component quant config. + + Pre-quantized formats (modelopt, svdquant, ...) require serialized scale + tensors the BF16 text-encoder checkpoint does not ship, so they are stripped + and the encoder stays in checkpoint precision. + """ + resolved = resolve_component_quant_config(quant_config, "text_encoder") + return resolve_encoder_quant_config(resolved) + + +# Subtrees of the HF Qwen3-VL text encoder that must stay in checkpoint +# precision: the vision tower (condition-image encoding is numerically +# sensitive) and the LM head (unused — the pipeline reads hidden states). +_TEXT_ENCODER_QUANT_EXCLUDED_PREFIXES = ("model.visual", "lm_head") + + +def _exclude_text_encoder_subtrees_from_quant( + quant_config: "QuantizationConfig | None", +) -> "QuantizationConfig | None": + """Keep the vision tower and LM head in checkpoint precision. + + Uses the config's standard ``ignored_layers`` mechanism, with substring + matching so the prefixes cover whole subtrees (same convention as the + DiT's ``_enable_pattern_ignored_layers``). + """ + if quant_config is None: + return None + if not hasattr(quant_config, "ignored_layers"): + logger.warning( + "Quantization config %s has no ignored_layers; the Qwen3-VL vision tower " + "and LM head cannot be excluded from quantization.", + type(quant_config).__name__, + ) + return quant_config + config = copy.copy(quant_config) + config.ignored_layers = [*config.ignored_layers, *_TEXT_ENCODER_QUANT_EXCLUDED_PREFIXES] + if hasattr(config, "ignored_layers_match_mode"): + config.ignored_layers_match_mode = "substring" + return config + + def get_qwen_image_21_pre_process_func( od_config: OmniDiffusionConfig, ): @@ -280,13 +330,41 @@ class QwenImage21Pipeline( ) # Unlike the 2.0 T2I pipeline, the vision tower is required: condition # images are encoded jointly with the prompt by the VLM. - self.text_encoder = from_pretrained_with_prefetch( - Qwen3VLForConditionalGeneration.from_pretrained, - model, - subfolder="text_encoder", - prefetch_list=qwen_subfolders, - local_files_only=local_files_only, - ).to(self.device) + text_encoder_quant_config = _resolve_text_encoder_quant_config(od_config.quantization_config) + if text_encoder_quant_config is None: + self.text_encoder = from_pretrained_with_prefetch( + Qwen3VLForConditionalGeneration.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=qwen_subfolders, + local_files_only=local_files_only, + ).to(self.device) + else: + # Quantized path: build the HF model on meta, swap its nn.Linear + # modules for vLLM quantizable linears, and stream the weights + # through the pipeline loader (same integration as Z-Image). The + # encoder-scoped quant config rides on a copied od_config so the + # shared create_transformers_model signature stays untouched. + text_encoder_config = AutoConfig.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) + encoder_od_config = dataclasses.replace( + od_config, + quantization_config=_exclude_text_encoder_subtrees_from_quant(text_encoder_quant_config), + ) + self.text_encoder = create_transformers_model( + AutoModelForImageTextToText, + encoder_od_config, + hf_config=text_encoder_config, + ).to(self.device) + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="text_encoder", + revision=None, + prefix="text_encoder.", + ) + ) self.vae = from_pretrained_with_prefetch( DistributedAutoencoderKLQwenImage21.from_pretrained, model, @@ -297,7 +375,7 @@ class QwenImage21Pipeline( transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImage21Transformer2DModel) self.transformer = QwenImage21Transformer2DModel( od_config=od_config, - quant_config=od_config.quantization_config, + quant_config=resolve_component_quant_config(od_config.quantization_config, "transformer"), **transformer_kwargs, ) self.processor = from_pretrained_with_prefetch( diff --git a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py index 07fc17259..7200cf7c7 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py @@ -348,6 +348,30 @@ class QwenImage21Rope(nn.Module): return torch.cat([self.freqs[0][frame_index], self.freqs[1][height_index], self.freqs[2][width_index]], dim=-1) +def _enable_pattern_ignored_layers(quant_config: QuantizationConfig | None) -> QuantizationConfig | None: + """Switch vLLM FP8-style `ignored_layers` to pattern (substring) matching. + + vLLM's ``Fp8Config`` matches ``ignored_layers`` against full layer prefixes + exactly, but the diffusion quantization docs and CLI describe them as name + patterns (e.g. ``img_mlp``) — under exact matching such patterns silently + skip nothing. Substring matching is a strict superset for exact full + prefixes, so existing exact-prefix configs keep working. + """ + if quant_config is None: + return quant_config + component_configs = getattr(quant_config, "component_configs", None) + if component_configs is not None: + configs = [*component_configs.values(), getattr(quant_config, "default_config", None)] + else: + configs = [quant_config] + for config in configs: + if config is None: + continue + if getattr(config, "ignored_layers", None) and hasattr(config, "ignored_layers_match_mode"): + config.ignored_layers_match_mode = "substring" + return quant_config + + class QwenImage21Attention(nn.Module): r"""Single-stream attention for Qwen-Image 2.1. @@ -736,7 +760,8 @@ class QwenImage21Transformer2DModel(CachedTransformer): self.inner_dim = num_attention_heads * attention_head_dim self.causal_condition = causal_condition self.causal_block = causal_block - self.quant_config = quant_config + self.quant_config = _enable_pattern_ignored_layers(quant_config) + quant_config = self.quant_config self.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(axes_dims_rope)) self.time_text_embed = QwenImage21TimestepProjEmbeddings(