[Feature]: FP8 Quantization Support for DiT

[lishunyang12]

FP8 Quantization Support for DiT (Z-Image & Qwen-Image)

Overview

This PR introduces online FP8 quantization for DiT models in vLLM-Omni, starting with Z-Image and Qwen-Image (text-to-image). Online FP8 converts BF16/FP16 weights to FP8 at model load time with dynamic activation scaling — no calibration dataset or pre-quantized checkpoint needed.

• Device compatibility: W8A8 on Ada/Hopper (SM 89+), weight-only via Marlin on Turing (SM 75+)
• Per-layer control: ignored_layers lets users keep sensitive layers in BF16

Supported Models

Model	Architecture	Recommendation	ignored_layers
Z-Image	S3-DiT (6B)	All layers quantized	None
Qwen-Image / Qwen-Image-2512	DiT	Skip sensitive layers	img_mlp

Changes

Quantization framework (vllm_omni/diffusion/quantization/)

• base.py — Abstract DiffusionQuantizationConfig wrapping vLLM's QuantizationConfig via quant_config_cls ClassVar
• fp8.py — DiffusionFp8Config with tight coupling to vLLM's Fp8Config (dynamic activation scaling, online weight conversion)
• __init__.py — Registry pattern, factory get_diffusion_quant_config(), and get_vllm_quant_config_for_layers() helper

Config integration (vllm_omni/diffusion/data.py)

• Added quantization: str | None and quantization_config fields to OmniDiffusionConfig
• __post_init__ handles dict/DictConfig/DiffusionQuantizationConfig input with Mapping ABC for OmegaConf compatibility

Model loader (vllm_omni/diffusion/model_loader/diffusers_loader.py)

• Added _process_weights_after_loading() to call quant_method.process_weights_after_loading() on quantized modules

Z-Image (vllm_omni/diffusion/models/z_image/)

• z_image_transformer.py — Added quant_config parameter to ZImageAttention, FeedForward, ZImageTransformerBlock, ZImageTransformer2DModel; passed to all vLLM parallel linear layers (QKVParallelLinear, RowParallelLinear, MergedColumnParallelLinear)
• pipeline_z_image.py — Extract and pass quant_config to transformer

Qwen-Image (vllm_omni/diffusion/models/qwen_image/)

• qwen_image_transformer.py — Added quant_config parameter to FeedForward, QwenImageCrossAttention, QwenImageTransformerBlock, QwenImageTransformer2DModel; passed to all vLLM parallel linear layers (QKVParallelLinear, RowParallelLinear, ColumnParallelLinear)
• pipeline_qwen_image.py — Extract and pass quant_config to transformer

Tests (tests/diffusion/quantization/)

• test_fp8_config.py — Unit tests covering config creation, vLLM config extraction, ignored_layers, dict non-mutation, conflicting method warnings, and end-to-end integration with OmniDiffusionConfig

Documentation (docs/)

• user_guide/diffusion/quantization/overview.md — Quantization methods overview
• user_guide/diffusion/quantization/fp8_online.md — Usage guide (Python API + CLI), parameter reference, per-model recommendations
• user_guide/diffusion_acceleration.md — Updated model support table with FP8 column
• .nav.yml — Added quantization section to docs navigation

Example (examples/offline_inference/text_to_image/text_to_image.py)

• Added --quantization fp8 and --ignored-layers CLI arguments

How to Use

from vllm_omni import Omni

# Z-Image: all layers quantized
omni = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="fp8")

# Qwen-Image: skip sensitive img_mlp layers
omni = Omni(
    model="Qwen/Qwen-Image-2512",
    quantization_config={
        "method": "fp8",
        "ignored_layers": ["img_mlp"],
    },
)

# CLI
python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization fp8

python text_to_image.py --model Qwen/Qwen-Image-2512 --quantization fp8 \
    --ignored-layers "img_mlp"

Benchmark (H200 x1)

Model	Task	Memory BF16	Memory FP8	Reduction	Time BF16	Time FP8	Speedup
Z-Image (6B)	Text-to-Image 1024x1024, 50 steps	19.15 GiB	13.56 GiB	29%	5.03s	3.83s	24%
Qwen-Image	Text-to-Image 1024x1024, 50 steps	53.74 GiB	41.20 GiB	23%	6.75s	5.26s	22%

Qwen-Image FP8 Ablation Study

An ablation over Qwen-Image-2512's 6 quantizable layer types (60 transformer blocks) found that img_mlp is the only FP8-sensitive layer. The image MLP processes denoising latents whose distribution shifts across timesteps; FP8's limited mantissa clips distribution tails, and errors compound over 60 residual blocks with no normalization correction.

Recommended config: quantize all layers except img_mlp (4 of 6 layer types in FP8, 2 stay BF16).

Visual Comparison

Z-Image — BF16 vs FP8 (seed=42, 1024x1024, 50 steps)

Prompt	BF16	FP8
"a cup of coffee on the table"
"a cat sitting on a windowsill at sunset"
"a mountain landscape with a lake reflection"

Not in Scope (Future PRs)

• Other DiT models (Flux, SD3, Wan2.2, GLM, LongCat, Ovis, Stable Audio) — same pattern applies; see #935
• Text encoder quantization — requires layer conversion like vLLM's Transformers backend
• Offline quantization (pre-quantized FP8 checkpoints, GGUF) — see #1285
• Bagel — uses HuggingFace transformers linear layers, not vLLM parallel layers

Related Issues

• RFC: FP8 Quantization Support for Diffusion Transformers #935
• RFC: Diffusion Models Quantization Support #896

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b7604ae37a

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Normalize DictConfig quantization_config before type check

If quantization_config is supplied via stage YAML or Omni(..., quantization_config={...}), OmegaConf turns that dict into a DictConfig. The isinstance(..., dict) guard then skips conversion, leaving od_config.quantization_config as a DictConfig. pipeline_z_image later calls get_vllm_quant_config_for_layers(od_config.quantization_config), which expects a DiffusionQuantizationConfig and will raise AttributeError on get_vllm_quant_config for a DictConfig. This breaks any run that uses per-method kwargs (e.g., activation_scheme). Consider accepting Mapping/DictConfig or converting with OmegaConf.to_container before the type check.

Useful? React with 👍 / 👎.

[Copilot]

Pull request overview

This pull request introduces online FP8 quantization support for the Z-Image diffusion transformer model in vLLM-Omni. The implementation enables dynamic activation scaling with BF16→FP8 weight conversion at runtime, requiring no pre-quantized checkpoints or calibration datasets.

Changes:

• Added quantization infrastructure with base class and FP8 implementation that wraps vLLM's Fp8Config
• Extended Z-Image transformer model to accept and propagate quant_config parameter through all linear layers
• Added quantization field to OmniDiffusionConfig with automatic conversion to quantization config objects in __post_init__
• Updated example script to support --quantization fp8 CLI argument with comprehensive unit tests

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
vllm_omni/diffusion/quantization/base.py	Base abstract class for diffusion quantization configs with vLLM config extraction interface
vllm_omni/diffusion/quantization/fp8.py	FP8-specific config wrapping vLLM's Fp8Config with dynamic activation scaling defaults
vllm_omni/diffusion/quantization/init.py	Factory functions and registry for creating quantization configs from string names
vllm_omni/diffusion/models/z_image/z_image_transformer.py	Added quant_config parameter to all transformer components and linear layer instantiations
vllm_omni/diffusion/models/z_image/pipeline_z_image.py	Integrated quantization config extraction and passing to transformer initialization
vllm_omni/diffusion/data.py	Added quantization fields to OmniDiffusionConfig with post_init conversion logic
examples/offline_inference/text_to_image/text_to_image.py	Added --quantization CLI argument with fp8 option and help text
tests/diffusion/quantization/test_fp8_config.py	Comprehensive unit tests for FP8 config creation, extraction, and error cases
tests/diffusion/quantization/init.py	Test module initialization file

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Using pop() on self.quantization_config when it's a dict will mutate the original dictionary. This could cause issues if the same config dict is reused elsewhere. Consider using dict.get() instead or creating a copy of the dict before popping.

Suggested fix: Replace quant_method = self.quantization_config.pop("method", self.quantization) with quant_method = self.quantization_config.get("method", self.quantization) and then pass a copy of the remaining kwargs.

Suggested change

	quant_method = self.quantization_config.pop("method", self.quantization)
	self.quantization_config = get_diffusion_quant_config(quant_method, **self.quantization_config)
	quant_method = self.quantization_config.get("method", self.quantization)
	quant_kwargs = {k: v for k, v in self.quantization_config.items() if k != "method"}
	self.quantization_config = get_diffusion_quant_config(quant_method, **quant_kwargs)

[Copilot]

The type annotation for quantization_config uses Any | None which is less specific than it could be. For consistency with cache_config (line 263), consider using a more specific type annotation.

Suggested change:

quantization_config: "DiffusionQuantizationConfig | dict[str, Any] | None" = None

Note: You'll need to add DiffusionQuantizationConfig to a TYPE_CHECKING import block to avoid circular imports, similar to how it's done in other files.

[ZJY0516]

Could you please try to add quant support for text_encoder?

[SamitHuang]

it's a good starting point. I think it worths a design doc and usage doc for future extension.

[SamitHuang]

to be general and consistent, i think it's better to add quant_config to other diffusion model implementation

[lishunyang12]

Sure, I will do. Bagel uses linear layer implementation from transformer lib which is not Compatible with vLLM built-in one. Except that, all Dit model can reuse this FP8 support.

[Isotr0py]

QQ: Do we need to call process_weights_after_loading at model loader for quantization?
https://github.com/vllm-project/vllm/blob/1bd47d6e5a57000e404af60838eadb86b6f457c0/vllm/model_executor/model_loader/utils.py#L81-L118

[lishunyang12]

Yes, i think so.

[Isotr0py]

Could you please try to add quant support for text_encoder?

Given that text_encoder is using AutoModel for model initialization, I think we will need some layer conversion like vLLM's Transformers backend, which may make this PR much more complicated.
https://github.com/vllm-project/vllm/blob/1bd47d6e5a57000e404af60838eadb86b6f457c0/vllm/model_executor/models/transformers/utils.py#L100-L137

[lishunyang12]

@david6666666 All DiT models(except Bagel) have been tested.

[Copilot]

Pull request overview

Copilot reviewed 40 out of 40 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The FeedForward import from diffusers is still used for glyph_projector and prior_projector (lines 664, 666) but these layers are not receiving the quant_config parameter. For consistency with the rest of the PR, these should either:

1. Use GlmImageFeedForward instead and receive quant_config, or
2. Have a documented reason why they should not be quantized.

If these projectors are part of the transformer and use linear layers, they should support FP8 quantization for maximum memory savings.

Suggested change

	import torch.nn.functional as F
	import torch.nn.functional as F
	# NOTE: We intentionally keep diffusers' FeedForward for small projector
	# modules (e.g., `glyph_projector` and `prior_projector`). These layers are
	# not wired into the FP8 quantization path on purpose: they are lightweight
	# projection heads whose memory footprint is small compared to the main
	# transformer blocks, and quantizing them can introduce undesirable numerical
	# effects for marginal memory savings. The main transformer feed-forward
	# layers instead use `GlmImageFeedForward` with `quant_config` for FP8
	# quantization.

[Copilot]

The dtype casting logic added here (lines 485-490) assumes that self.pos_embed.proj.weight always exists and has a valid dtype. However, with FP8 quantization, the weight might be in FP8 format while we want to cast inputs to the original dtype (BF16/FP16). Consider whether model_dtype should be obtained from a more reliable source, such as self.dtype or the config, rather than from potentially quantized weights.

Additionally, this casting logic is only added to SD3 but not to other models. If this is necessary for FP8 quantization to work correctly, it should be consistently applied across all models that support FP8.

Suggested change

	model_dtype = self.pos_embed.proj.weight.dtype
	hidden_states = hidden_states.to(dtype=model_dtype)
	encoder_hidden_states = encoder_hidden_states.to(dtype=model_dtype)
	pooled_projections = pooled_projections.to(dtype=model_dtype)
	model_dtype = getattr(self, "dtype", None)
	if model_dtype is None:
	model_dtype = hidden_states.dtype
	if hidden_states.dtype is not model_dtype:
	hidden_states = hidden_states.to(dtype=model_dtype)
	if encoder_hidden_states is not None and encoder_hidden_states.dtype is not model_dtype:
	encoder_hidden_states = encoder_hidden_states.to(dtype=model_dtype)
	if pooled_projections is not None and pooled_projections.dtype is not model_dtype:
	pooled_projections = pooled_projections.to(dtype=model_dtype)

[david6666666]

I think this online quantization method is relatively niche, diffusion is sensitive to quantization accuracy; most people probably use GGUF and native FP8 weights more often. I understand that our online quantization method will have a significant impact on accuracy. If possible, could you provide a few typical models to demonstrate the accuracy comparison?
@Isotr0py @SamitHuang @ZJY0516 @wtomin @lishunyang12 WDYT

[ZJY0516]

I'm not sure the differnce between online fp8 quant and native fp8 weights

[Isotr0py]

I understand that our online quantization method will have a significant impact on accuracy

I'm not sure whether the existing diffusion fp8 checkpoints use dynamic fp8 (per channel) or static fp8 (per token). If it uses dynamic FP8, there should be no difference between online and offline quantization.

For static FP8 offline quantization, it relatively rely on the calibration dataset, but seems that there is no tools to calibrate diffusion models during quantization yet?

[lishunyang12]

I understand that our online quantization method will have a significant impact on accuracy

I'm not sure whether the existing diffusion fp8 checkpoints use dynamic fp8 (per channel) or static fp8 (per token). If it uses dynamic FP8, there should be no difference between online and offline quantization.

For static FP8 offline quantization, it relatively rely on the calibration dataset, but seems that there is no tools to calibrate diffusion models during quantization yet?

As been said in https://github.com/vllm-project/vllm-omni/pull/1034/changes#diff-c1202f8fbeb964dc0a65b3e9ef72f605c70ea8ab121001e5e8d101f200f5d4caR10-R29, the current online quantization uses dynamic activation scaling (no calibration dataset needed) and online weight quantization from BF16/FP16 checkpoints. As for the precision loss concern,I will share some results for eyeballing the precision before and after quantization for somemodels. @david6666666 @ZJY0516 @Isotr0py

[lishunyang12]

I think this online quantization method is relatively niche, diffusion is sensitive to quantization accuracy; most people probably use GGUF and native FP8 weights more often. I understand that our online quantization method will have a significant impact on accuracy. If possible, could you provide a few typical models to demonstrate the accuracy comparison? @Isotr0py @SamitHuang @ZJY0516 @wtomin @lishunyang12 WDYT

I dont think online quantization method is niche.

[lishunyang12]

Text to Image

Z-Image:

Prompt: A breathtaking twilight scene atop a floating archipelago of crystalline islands suspended in an endless nebula-drenched sky, where waterfalls of liquid starlight cascade from the edges of each island into the cosmic abyss below. The largest island hosts an ancient, overgrown observatory crafted from pearlescent white stone and living silverwood trees whose bioluminescent leaves pulse with soft cerulean and violet light. At the observatory's center stands a solitary figure—a young astronomer in flowing robes woven from woven moonlight and deep-space fabric, their face illuminated by the glow of a holographic star chart hovering above an altar of polished obsidian. The sky swirls with the birth of a new galaxy: ribbons of magenta and gold gas clouds twist around nascent stars, while distant supernovae bloom like cosmic flowers in the far reaches of the void. Below, the abyss reveals faint echoes of forgotten civilizations—ghostly silhouettes of submerged cities shimmering in layers of atmospheric haze. Cinematic volumetric lighting cuts through the scene as twin moons—one copper-hued, one opalescent—rise on opposite horizons, casting long, intersecting shadows across moss-covered ruins and crystalline flora that refract light into prismatic halos. Hyper-detailed, photorealistic rendering with the atmospheric depth of Roger Deakins' cinematography, the architectural grandeur of Zaha Hadid, and the cosmic wonder of James Jean's illustrations. Shot on a mythical 150mm lens with shallow depth of field, 8K resolution, Unreal Engine 5 realism, subsurface scattering on organic elements, and ray-traced reflections dancing across every water droplet in the starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation at the edge of existence.

BF16:

FP8:

Text to Video

Wan2.2

Prompt: Two expressive anthropomorphic cats duel in a jazz-age boxing ring under a single swinging spotlight: a lanky ginger tabby in emerald velour shorts and oversized pearl-button gloves executes a lightning-fast triple-jab combo, fur rippling with motion blur as sweat droplets fly in parabolic arcs. His opponent—a stocky tuxedo cat in crimson satin trunks—ducks under the assault with a playful grin, then counters with a spinning back-paw hook that sends a shimmering shockwave through the air (visible as concentric heat-haze rings distorting the background crowd). Mid-sequence (4-6 seconds): the spotlight cable snaps, causing the beam to swing wildly in a 270° arc—casting rapidly shifting chiaroscuro shadows that race across the ring canvas, illuminating floating chalk dust and individual whiskers twitching with exertion. Final seconds: both cats leap simultaneously for a mid-air clash, gloves colliding in a burst of golden sparks that hang suspended for 3 frames before drifting downward like fireflies. Camera executes a dynamic dolly-zoom during the collision—background stretching while foreground compresses—enhancing the impact's visceral weight. Photorealistic fur simulation with subsurface scattering on ear membranes backlit by sparks, motion blur graded by velocity (sharp paws → smeared glove trails), 16fps cinematic cadence with squash-and-stretch animation principles applied to feline musculature.

BF16:
https://github.com/user-attachments/assets/0091c67b-b432-45e1-b5fb-d4e92be5919d

FP8
https://github.com/user-attachments/assets/2f2ad4a3-0fae-4bfb-af33-9cdaeeaea68e

Image to Image - Bugs!!!

Qwen Image Edit

Original Image:

Prompt: The bear should have its name puppy on his belly, Put a Santa hat on the puppy, change the background to night with a bright moon.

BF16:

FP8:

Image to Video

Wan2.2

Original Image:

Prompt: A fluffy teddy bear joyfully tumbling through a snowy night landscape, its round belly prominently displaying the stitched name 'PUPPY' in playful red thread. A crooked Santa hat bounces precariously on its head as it rolls, one paw reaching up to keep it from falling off. Behind the action, a luminous full moon casts silvery beams through swirling snowflakes, illuminating the bear's mid-tumble motion with dramatic shadows and highlights. Pine trees blur in the background as the bear cartwheels across a frost-dusted hillside under a star-speckled indigo sky
BF16:
https://github.com/user-attachments/assets/c5eeb927-152e-4383-b3db-6ae4a6f4c163
FP8:
https://github.com/user-attachments/assets/d7b66c4b-c449-4382-9c61-3fd4c7e91369

[lishunyang12]

This PR is ready.

[david6666666]

please resolve conflicts

[lishunyang12]

please resolve conflicts

Resolved.

[david6666666]

@Isotr0py @ZJY0516 please review 🙏 thx

[ZJY0516]

LGTM now

[david6666666]

LGTM