bahadirakdemir/Z-Image-Turbo-tensorrt-onnx-fp8

Z-Image-Turbo — TensorRT FP8 ONNX (DiT only)

This repo ships FP8-quantized DiT engines for Tongyi-MAI/Z-Image-Turbo plus the matching BF16 text encoder. The DiT is the only part quantized; encoder and VAE stay BF16 per industry pattern (FLUX-FP8, SD3.5-FP8).

For the full BF16 baseline (1.0× speed reference), see bahadirakdemir/Z-Image-Turbo-onnx-bf16.

Files

File	Pair file	Size	Notes
qwen3_text_encoder.onnx	.onnx.data	7.85 GB	Qwen3 text encoder (BF16, same as the BF16 repo). Output is the penultimate hidden state.
zimage_dit_1024x512_fp8.onnx	.onnx.data	6.17 GB	DiT @ 1024×512, FP8-quantized via NVIDIA ModelOpt. Half the weight footprint of the BF16 variant.
zimage_dit_512x512_fp8.onnx	.onnx.data	6.17 GB	DiT @ 512×512, FP8-quantized.

ONNX opset 18. External weights are in the .onnx.data sidecars.

⚠ Important: FP8 ONNX is TensorRT-specific

The FP8 DiT graphs use TRT_FP8DequantizeLinear / TRT_FP8QuantizeLinear ops from the trt opset domain. These are NOT standard ONNX operators. The FP8 ONNX files are only consumable by TensorRT, not by ONNX Runtime or any generic ONNX backend. The BF16 encoder ONNX is standard ONNX and runs anywhere.

Performance

End-to-end latency and peak GPU memory on a single NVIDIA GB10, 8 inference steps, guidance_scale=0.0, batch 1. All five backends ran the same upstream Tongyi-MAI/Z-Image-Turbo at HF revision f332072a. n = 10 prompts per resolution, mean shown (run-to-run within 1 %). Diffusers-Server and vllm-omni latencies include local-socket HTTP round-trip (50 ms). TRT rows measured standalone (HF pipeline dropped from GPU after engine load). Memory captured via nvidia-smi --query-compute-apps because GB10 doesn't expose the GPU-wide memory.used gauge.

Backend (loaded engine set)	1024×512 latency	512×512 latency	Peak GPU (1024×512 / 512×512)
HF BF16 (PyTorch + ZImagePipeline)	5.84 s	2.64 s	22.26 / 22.26 GB
diffusers-server BF16 (HTTP)	7.32 s	3.47 s	22.26 / 22.26 GB
vllm-omni BF16 (HTTP, TORCH_SDPA)	7.14 s	3.73 s	21.83 / 21.83 GB
TRT BF16, both engines loaded (sibling repo)	4.95 s	2.14 s	34.20 / 34.20 GB
TRT BF16, single engine (sibling repo)	5.10 s	2.15 s	22.49 / 20.95 GB
TRT FP8, both engines loaded (this repo)	2.30 s	1.12 s	22.07 / 22.07 GB
TRT FP8, single engine (this repo)	2.34 s	1.15 s	16.23 / 15.08 GB
TRT NVFP4-W4A8, single engine (experimental)	4.84 s	2.78 s	14.20 / 13.07 GB

Speedup vs each baseline (1024×512 / 512×512):

Backend	vs HF	vs Diffusers-Server	vs vllm-omni	vs TRT BF16
TRT FP8 (this repo)	2.54× / 2.36×	3.18× / 3.10×	3.10× / 3.33×	2.15× / 1.91×

Memory. With a single DiT engine loaded (the production pattern — load only the resolution you serve), TRT FP8 peaks at 15–16 GB, less than every BF16 backend tested, while running 2.5–3× faster. Loading both DiTs simultaneously costs the second engine's weights (~6 GB FP8, ~12 GB BF16) — worth it only if you switch resolutions per-request and the reload cost is unacceptable. See "Reducing memory further" below for additional options.

Reducing memory further

• Load engines per-resolution — the 16 GB row above. The cleanest production pattern; sacrifices ~50 ms when you have to swap engines between resolutions (deserialize a different .plan from disk).
• Single dynamic-shape engine (not provided here) — re-export with an optimization profile covering both resolutions, build one engine instead of two. Saves ~6 GB FP8 / ~12 GB BF16, costs roughly 10–15 % latency.
• Weight streaming (TensorRT 10+) — pass BuilderFlag.WEIGHT_STREAMING at build time and call engine.weight_streaming_budget = .... Keeps weights on host and streams per-layer; can roughly halve resident GPU memory at the cost of measurable latency. Worth exploring on ≤ 16 GB cards.
• Drop the BF16 encoder if running multiple worker processes — only one worker needs to hold the encoder; the rest can RPC to it. Saves the 7.85 GB encoder per replica.
• NVFP4-W4A8 DiT (experimental, last row in the table) — weights in NVFP4 (~3.87 GB plan vs FP8's 6.17 GB), activations in FP8. Peak GPU drops to 13–14 GB (smallest of any backend). On TRT 10.16.1 the compute path is FP8×FP8 with FP4→FP8 dequant per matmul, so latency is ~2× slower than FP8 (back to BF16-PyTorch territory). Worthwhile only on memory-constrained cards (≤ 16 GB VRAM) where FP8 doesn't quite fit. A native NVFP4 tensor-core path will land in TRT > 10.16 and should restore the speed advantage.

DiT-only step rate (the part FP8 quantization actually accelerates):

Backend	1024×512	512×512
HF BF16	1.60 it/s	3.55 it/s
TRT BF16	1.74 it/s	4.10 it/s
TRT FP8	4.10 it/s	8.65 it/s

Image quality: visually equivalent to the BF16 baseline on photoreal, instruction-following, and EN+ZH bilingual text-rendering prompts. Pixel-level rmse drift vs BF16 reference: ~30–35 on a 0–255 scale (text-rendering prompts can hit ~50–70 due to single-bit-flip sensitivity of fine glyph strokes, but text remains correctly rendered).

Quantization recipe

• Tool: NVIDIA TensorRT Model Optimizer (ModelOpt) 0.44.0.
• Config: mtq.FP8_DEFAULT_CFG (E4M3, per-tensor max calibration).
• Excluded modules (kept BF16, model declares them precision-sensitive via _skip_layerwise_casting_patterns):
  • *t_embedder* — timestep embedding (TimestepEmbedder, 256→1024→3840).
  • *cap_embedder* — caption-feature projection (RMSNorm + Linear, 2560→3840).
• Calibration data: 288 DiT forward inputs captured from the HF BF16 pipeline running 16 fresh prompts × 2 resolutions × 9 sigma points (one full sweep through the FlowMatchEulerDiscreteScheduler at num_inference_steps=9).
• Total quantized modules: 1299 (FP8 Linear projections + adaLN modulations). 12 modules excluded per the rules above.

Engine input / output schemas

Encoder (qwen3_text_encoder.onnx)

Identical to the BF16 repo:

Tensor	Direction	Shape	Dtype
input_ids	input	[1, 512]	int64
attention_mask	input	[1, 512]	int64
penultimate_hidden_state	output	[1, 512, 2560]	bfloat16

DiT FP8 (zimage_dit_<resolution>_fp8.onnx)

Same schema as the BF16 variant (FP8 Q/DQ nodes are internal to the graph):

Tensor	Direction	Shape	Dtype	Notes
latent	input	[1, 16, H_lat, W_lat]	bfloat16	H_lat=H/8, W_lat=W/8.
t	input	[1]	float32	Scheduler sigma.
cap_feats	input	[1, 1024, 2560]	bfloat16	Caller pads to T_cap=1024.
cap_attn_mask	input	[1, 1024]	bool	True for valid caption tokens.
freqs_cis_x	input	[1, T_x, 64, 2]	float32	Pre-computed cos/sin RoPE.
freqs_cis_cap	input	[1, 1024, 64, 2]	float32	Pre-computed cos/sin RoPE.
noise_pred	output	[1, 16, H_lat, W_lat]	bfloat16

Image-token counts: T_x = (H/8) * (W/8) / 4 after patch_size=2. For 1024×512: 2048; for 512×512: 1024.

Important caveats

(See the BF16 repo README for the full version; summarized here.)

1. Real cos/sin RoPE. The export rewrites the upstream complex-tensor RoPE to real cos/sin pairs. Pass freqs_cis_* as [1, T, 64, 2] tensors; do NOT pass complex tensors.
2. Caller applies the chat template. Run tokenizer.apply_chat_template([{role: user, content: prompt}], add_generation_prompt=True, enable_thinking=True) before tokenizing. The encoder engine has no chat template inside.
3. Penultimate hidden state. The encoder returns hidden_states[-2] (matches pipeline_z_image.py:236). It is NOT last_hidden_state.
4. Static T_cap=1024. Caption tokens (after attention_mask trimming) must be ≤ 1024 for the static-shape DiT. Raised from 128 to fit longer compiled image prompts (~430–480 tokens worst-case from upstream orchestrators); ~2× headroom for future prompt growth. Trade-off: DiT step is ~30 % slower than the 128-cap variant at 1024×512 (FP8) due to the longer single-stream attention sequence. See the matching BF16 repo for the same shape.

Building TensorRT engines

import tensorrt as trt

def build_fp8(onnx_path, plan_path):
    logger = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 32 << 30)
    # FP8 with BF16 fallback for excluded layers (D-012: t_embedder, cap_embedder)
    config.set_flag(trt.BuilderFlag.BF16)
    config.set_flag(trt.BuilderFlag.FP8)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    assert parser.parse_from_file(onnx_path)
    serialized = builder.build_serialized_network(network, config)
    with open(plan_path, "wb") as f: f.write(serialized)

build_fp8("qwen3_text_encoder.onnx",       "qwen3_text_encoder.plan")        # BF16 only (no FP8 in encoder)
build_fp8("zimage_dit_1024x512_fp8.onnx",  "zimage_dit_1024x512_fp8.plan")
build_fp8("zimage_dit_512x512_fp8.onnx",   "zimage_dit_512x512_fp8.plan")

Plan sizes on GB10: encoder 7.85 GB, each FP8 DiT **6.17 GB** (vs ~12.31 GB BF16).

Build time on GB10: encoder ~30 s, each DiT FP8 ~65–70 s.

Hardware requirements

• Native FP8 GPU: Hopper (H100/H200, sm_90), Blackwell (GB10, B100, B200; sm_100/120/121).
• TensorRT: ≥ 10.0 with FP8 support. Verified on TRT 10.16.1.
• Driver / CUDA: anything that pairs with your TRT version. Verified on driver 595.58.03, CUDA 13.0.88.
• GPU memory: at least 24 GB for the loaded FP8 engines plus working memory; 48 GB recommended.
• Ada and older: NOT supported for the FP8 path (no native FP8). Use the BF16 sibling repo instead.

License & attribution

Apache-2.0, inherited from the upstream model:

• Upstream model: Tongyi-MAI/Z-Image-Turbo © Alibaba Z-Image Team. See their LICENSE and model card.
• Papers: arXiv 2511.22699, 2511.22677, 2511.13649.
• Quantization: NVIDIA TensorRT Model Optimizer (Apache-2.0).
• This export: pinned at upstream revision f332072aa78be7aecdf3ee76d5c247082da564a6 (HF, 2026-01-30).

Please cite the original Z-Image team's work when using.