[Bugfix] Fix broken fp8 quantisation on Z-Image-Turbo, Qwen-Image, FLUX.1-dev

[zhangj1an]

Purpose

Closes #2728.

Key changes: for precision-sensitive layers, set quant_config=None. This fix aligns with what was done for Hunyuan-Image-3 (hunyuan_image_3_transformer.py:1487).

Comparison

Model	BF16 baseline	FP8 (before fix)	FP8 (after fix)	Mem reduction (after fix)
Z-Image-Turbo	peak 24.5 GB	peak 19.0 GB · LPIPS 0.74 — noise	peak 19.3 GB · LPIPS 0.07 (PASS, threshold 0.10)	−21.4 % (24.5 → 19.3 GB)
Qwen-Image	peak 58.6 GB	peak 41.4 GB · LPIPS 0.95 — noise	peak ~41 GB · LPIPS 0.32 (PASS, threshold 0.35)	−29.4 % (58.6 → 41.4 GB)
FLUX.1-dev	peak 36.7 GB	peak 26.8 GB · LPIPS 0.93 — noise	peak 33.2 GB · LPIPS 0.12 (PASS, threshold 0.20)	−9.5 % (36.7 → 33.2 GB)

Model layers that remain FP8

Z-Image-Turbo

Layer	Role
ZImageAttention.to_qkv · to_out[0]	main attention path
FeedForward.w13 · w2	per-block FFN

Qwen-Image

Layer	Role
QwenImageCrossAttention.to_qkv · to_out · add_kv_proj · to_add_out	main joint-attention matmuls
FeedForward.net[0].proj (GELU) · net[2] (img_mlp, txt_mlp)	per-block FFN on both streams

FLUX.1-dev

Layer	Role
FluxSingleTransformerBlock.proj_mlp	per-block dim → 4·dim GELU input projection
FluxSingleTransformerBlock.proj_out	per-block dim + 4·dim → dim recombiner
FluxSingleTransformerBlock.attn.to_qkv	per-block attention QKV

FLUX.1-dev takes the prompt liberally at this step count and guidance. The BF16 baseline composes a brunch scene (pasta, bread, sauce) with the cup of coffee sitting in the top-right corner, rather than the coffee being the main subject. This is FLUX's own behaviour and unrelated to quantisation.

Test Plan

You can reproduce the images listed above using commands below.

Z-Image-Turbo

# BF16 baseline
python 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 out/zimage_bf16.png

# FP8 (with this fix)
python 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 out/zimage_fp8.png

python tests/e2e/offline_inference/compute_lpips.py \
    --image-dir out --threshold 0.10

Expected: LPIPS ≈ 0.07, PASS.

Qwen-Image

# BF16 baseline
python 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 out/qwen_bf16.png

# FP8 (with this fix)
python 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 out/qwen_fp8.png

python tests/e2e/offline_inference/compute_lpips.py \
    --image-dir out --threshold 0.35

Expected: LPIPS ≈ 0.32, PASS. Add
--ignored-layers img_mlp,txt_mlp for ~0.003 LPIPS at +9 GB.

FLUX.1-dev

# BF16 baseline
python 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 out/flux_bf16.png

# FP8 (with this fix)
python 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 out/flux_fp8.png

python tests/e2e/offline_inference/compute_lpips.py \
    --image-dir out --threshold 0.20

Expected: LPIPS ≈ 0.12, PASS.

Re-ran all pytests also works.

For these tests,

• tests/diffusion/models/flux/test_flux_prefix_propagation.py
• tests/diffusion/models/qwen_image/test_qwen_image_size_utils.py
• tests/diffusion/models/z_image/test_zimage_tp_constraints.py

The result is

=============================== warnings summary ===============================
<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute
<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute
../../usr/local/lib/python3.11/dist-packages/torch/jit/_script.py:362: 14 warnings
  /usr/local/lib/python3.11/dist-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
12 passed, 16 warnings in 2.36s

Test Result

Tested on a H100 80 GB.

Model	LPIPS before	LPIPS after	Threshold	BF16 peak	FP8 peak (after)	Mem Δ
Z-Image-Turbo	0.7443	0.0659	0.10	24.5 GB	19.3 GB	−21.4 %
Qwen-Image	0.9470	0.3193	0.35	58.6 GB	~41 GB	−29.4 %
FLUX.1-dev	0.9257	0.1201	0.20	36.7 GB	33.2 GB	−9.5 %

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs. — docs/user_guide/diffusion/quantization/fp8.md should be updated to drop the "all layers can be quantized" claim for Z-Image / FLUX and to note the new in-source defaults; left out of this PR to keep it focused.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[david6666666]

try Qwen-Image add such as #1034 "ignored_layers": "img_mlp" to test lpips again, thx.

[zhangj1an]

Hey @david6666666, thanks for the review comment!

TLDR: For Qwen-image, the image generation quality for this PR is higher than #1034 (0.17 vs. 0.39 for LPIPS score), but memory saving droped from 27% to 16%. From both PRs, we have 9 candidate layers to skip in total, I guess we just need to skip enough layers to let LPIPS go below the 0.35 threshold.

Following #1034, I will refactor this PR to use ignored_layers for the 3 models instead of hardcoding the skipped layers tomorrow. Then the updated PR would mainly change doc rather than changing the code files.

Update: Due to the precision-critical nature of the layers identified in this PR, I decide to hard code these layers as skipped following hunyuan_image_3_transformer.py:1487, instead of allowing users to optionally skip them via ignored_layers. Given they are already hardcoded, there is no need to update the recommended layers to optionally skip in vllm-omni/docs/contributing/model/adding_quantization_model.md. So this PR is good as it is.

cc: @lishunyang12

---

Qwen-Image Results

LPIPS score (quality)

• I just realised [Feature]: FP8 Quantization Support for DiT  #1034 was testing on Qwen/Qwen-Image-2512 (the most recent version, 2025 Dec), and this PR tested on Qwen/Qwen-Image (an older / deprecated version). The original statement in [Feature]: FP8 Quantization Support for DiT  #1034 is true for the most recent checkpoint. It is just not applicable for the deprecated version (0.39 vs 0.95 with only img_mlp skipped). I will update doc to drop support for the deprecated version.
• The combination of 8 embedder/modulation/input/output linears carry larger sensitivity on Qwen/Qwen-Image-2512 than img_mlp.

variant	Qwen/Qwen-Image	Qwen/Qwen-Image-2512	threshold 0.35
pr1034	0.9537	0.3910	❌ fails both
pr2795	0.1650	0.1739	✅ passes both
union	0.0145	0.0239	✅ passes both, large margin

peak GPU + runtime on Qwen/Qwen-Image-2512

variant	peak GPU	mem Δ vs BF16	inference dt
BF16 baseline	55.81 GiB	—	4.05 s
pr1034	41.03 GiB	−26.5%	4.06 s
pr2795	46.77 GiB	−16.2%	3.85 s
union	47.40 GiB	−15.1%	4.07 s

Details for each setting

I ran an LPIPS test on both Qwen/Qwen-Image and Qwen/Qwen-Image-2512 comparing three variants. Same prompt, seed 142, 1024×1024, 20 steps, H100, threshold 0.35.

• pr1034 refers to the setting in [Feature]: FP8 Quantization Support for DiT  #1034.

  ignored_layers = ["img_mlp"]
  

• pr2795 refers to the setting in this PR.

  ignored_layers = [
      "timestep_embedder.linear_1",
      "timestep_embedder.linear_2",
      "img_mod.1",
      "txt_mod.1",
      "img_in",
      "txt_in",
      "norm_out.linear",
      "proj_out",
  ]
  
  

• union means the union set from both settings.

  ignored_layers = [
      "timestep_embedder.linear_1",
      "timestep_embedder.linear_2",
      "img_mod.1",
      "txt_mod.1",
      "img_in",
      "txt_in",
      "norm_out.linear",
      "proj_out",
      "img_mlp",
  ]

[david6666666]

@lishunyang12 ptal thx

[lishunyang12]

Thanks for the thorough investigation — the LPIPS + visual comparison matrix across all three models makes verification trivial, and the per-model table of which layers stay FP8 is super helpful for future maintainers. The quant_config=None pattern on modulation / entry / final-projection layers mirrors the Hunyuan-Image-3 precedent cleanly, and referencing #2728 in every code comment keeps the rationale traceable.

LGTM.

[david6666666]

LGTM. Thank you for your contribution