Add FP8 online qunatization support for Qwen-Image-2.1 (#7792)

Signed-off-by: Isotr0py <mozf@inferact.ai>

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)      |       ❌        |        ❌         |

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 | |

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
 

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(

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(