Skip to content

Commit 9d19ec6

Browse files
csgooglecsgoogle
committed
Optimize Z-Image-Turbo serving pipeline and attention layouts
- Pre-split text_encoder and transformer during initialization to avoid host-side splitting CPU overhead. - Fully tensorize ZImageTransformer2DModel __call__ forward path using statically-padded 3D stacked sequences, bypassing cross-device all-reduce/all-gather collectives and loops. - Pass 4D transposed layouts directly to NNXAttentionOp to bypass unflatten/flatten transposes. - Maintain RMSNorm outputs in float32 and pass them directly into _apply_rope to enable seamless on-device loop/kernel fusion and avoid float32 <-> bfloat16 casting loops. - Inline latent generation logic inside pipeline call, removing redundant JIT overhead.
1 parent bda41e1 commit 9d19ec6

22 files changed

Lines changed: 2491 additions & 226 deletions

‎.agents/plugins/tpu-catalyst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
/mnt/disks/data/tpu-catalyst

‎src/maxdiffusion/checkpointing/checkpointing_utils.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
STABLE_DIFFUSION_XL_CHECKPOINT = "STABLE_DIFUSSION_XL_CHECKPOINT"
3737
FLUX_CHECKPOINT = "FLUX_CHECKPOINT"
3838
WAN_CHECKPOINT = "WAN_CHECKPOINT"
39+
Z_IMAGE_CHECKPOINT = "Z_IMAGE_CHECKPOINT"
3940

4041

4142
def create_orbax_checkpoint_manager(
@@ -62,6 +63,16 @@ def create_orbax_checkpoint_manager(
6263
item_handlers = None
6364
if checkpoint_type == FLUX_CHECKPOINT:
6465
item_names = ("flux_state", "flux_config", "vae_state", "vae_config", "scheduler", "scheduler_config")
66+
elif checkpoint_type == Z_IMAGE_CHECKPOINT:
67+
# Only `transformer_state` is trainable; the VAE and the Qwen3 text encoder
68+
# are frozen and kept as separate non-trainable items.
69+
item_names = ("transformer_state", "vae_state", "text_encoder_state", "z_image_config")
70+
item_handlers = {
71+
"z_image_config": ocp.JsonCheckpointHandler(),
72+
"transformer_state": ocp.StandardCheckpointHandler(),
73+
"vae_state": ocp.StandardCheckpointHandler(),
74+
"text_encoder_state": ocp.StandardCheckpointHandler(),
75+
}
6576
elif checkpoint_type == WAN_CHECKPOINT:
6677
item_names = ("low_noise_transformer_state", "high_noise_transformer_state", "wan_state", "wan_config")
6778
item_handlers = {

‎src/maxdiffusion/checkpointing/z_image_checkpointer.py‎

Lines changed: 414 additions & 0 deletions
Large diffs are not rendered by default.

‎src/maxdiffusion/common_types.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
WAN2_2 = "wan2.2"
5454
LTX2_VIDEO = "ltx2_video"
5555
LTX2_3 = "ltx2.3"
56+
Z_IMAGE = "z_image"
5657

5758
WAN_MODEL = WAN2_1
5859

‎src/maxdiffusion/configs/base_zimage.yml‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# Official Z-Image inference defaults. Z-Image and Z-Image-Turbo share the
16+
# denoiser; only the recommended denoising-step count differs (see
17+
# base_zimage_turbo.yml).
18+
19+
run_name: 'zimage'
20+
21+
# If true save config to GCS in {base_output_directory}/{run_name}/
22+
save_config_to_gcs: False
23+
24+
pretrained_model_name_or_path: 'Tongyi-MAI/Z-Image'
25+
model_name: z_image
26+
model_type: 'T2I'
27+
28+
unet_checkpoint: ''
29+
# This will convert the weights to this dtype.
30+
# When running inference on TPUv5e, use weights_dtype: 'bfloat16'
31+
weights_dtype: 'bfloat16'
32+
# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype)
33+
activations_dtype: 'bfloat16'
34+
35+
# Maximum sequence length for the text encoder
36+
max_sequence_length: 512
37+
38+
# Attention
39+
attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
40+
# Maxdiffusion has 2 types of attention sharding strategies:
41+
# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
42+
# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is
43+
# sharded in cross attention q.
44+
attention_sharding_uniform: True
45+
# Empty dict lets the kernel pick defaults. Example v6e override:
46+
# flash_block_sizes: {
47+
# "block_q" : 1024,
48+
# "block_kv_compute" : 1024,
49+
# "block_kv" : 1024,
50+
# "block_q_dkv" : 1024,
51+
# "block_kv_dkv" : 1024,
52+
# "block_kv_dkv_compute" : 1024,
53+
# "block_q_dq" : 1024,
54+
# "block_kv_dq" : 1024,
55+
# "use_fused_bwd_kernel": False,
56+
# }
57+
flash_block_sizes: {}
58+
59+
# Output directory
60+
# Create a GCS bucket, e.g. my-maxdiffusion-outputs and set this to "gs://my-maxdiffusion-outputs/"
61+
base_output_directory: ""
62+
output_dir: "/mnt/disks/data/tmp/maxdiffusion"
63+
# Local path for the generated image. If empty, an image named
64+
# {run_name}_output_{seed}.png is written to the working directory.
65+
output_file: "/mnt/disks/data/tmp/zimage.png"
66+
67+
# Hardware
68+
hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
69+
skip_jax_distributed_system: True
70+
71+
# Parallelism
72+
mesh_axes: ['data', 'fsdp', 'context', 'tensor']
73+
# Z-Image runs sequence parallel (activations sharded over `context`) with
74+
# data parallelism; the tensor axis is left at 1. `attention_sharding_uniform`
75+
# adds the matching q/kv sequence rules on top of these.
76+
logical_axis_rules: [
77+
['batch', ['data', 'fsdp']],
78+
['activation_batch', ['data', 'fsdp']],
79+
['activation_length', 'context'],
80+
['activation_heads', 'tensor'],
81+
['heads', 'tensor'],
82+
['mlp', 'tensor'],
83+
['embed', 'fsdp'],
84+
['norm', 'tensor'],
85+
['out_channels', 'tensor'],
86+
]
87+
data_sharding: [['data', 'fsdp', 'context', 'tensor']]
88+
89+
# One axis for each parallelism type may hold a placeholder (-1)
90+
# value to auto-shard based on available slices and devices.
91+
#
92+
# Measured on v6e-8, Turbo, 1024x1024, 9 steps, 8 images (per_device_batch_size 1.0):
93+
# data=8 2.45s 0.307s/image <- default
94+
# data=4 ctx=2 2.60s 0.325s/image
95+
# data=2 ctx=4 3.13s 0.391s/image
96+
# ctx=8 4.15s 0.519s/image
97+
# fsdp=8 9.43s 1.179s/image
98+
# The 12B of weights fit in HBM, so replicating them and giving each chip a
99+
# whole image beats every sharded layout: no collectives at all. fsdp re-gathers
100+
# every weight on all 9 denoising steps with no backward pass to amortize it.
101+
#
102+
# For single-image *latency* instead of throughput, generate one image
103+
# (per_device_batch_size: 0.125) with ici_data_parallelism: 1 and
104+
# ici_context_parallelism: -1, which splits one image's sequence 8 ways: 0.587s.
105+
dcn_data_parallelism: 1
106+
dcn_fsdp_parallelism: 1
107+
dcn_context_parallelism: 1
108+
dcn_tensor_parallelism: 1
109+
ici_data_parallelism: -1 # recommended ICI axis to be auto-sharded
110+
ici_fsdp_parallelism: 1
111+
ici_context_parallelism: 1
112+
ici_tensor_parallelism: 1
113+
114+
allow_split_physical_axes: False
115+
116+
# Generation parameters
117+
prompt: "A red fox reading a book in a quiet library, cinematic light"
118+
height: 1024
119+
width: 1024
120+
num_inference_steps: 50
121+
# The released Z-Image configuration runs without classifier free guidance.
122+
guidance_scale: 0.0
123+
seed: 42
124+
# Images per device. Fractional values are allowed: on 8 devices,
125+
# 0.125 generates a single image.
126+
per_device_batch_size: 1.0
127+
# Latents decoded per VAE call. The decoder is replicated and its
128+
# activations are full-resolution, so raising this trades memory for speed.
129+
vae_decode_chunk: 1
130+
131+
# Profiling
132+
# generate_zimage runs an un-profiled warmup pass (compile), then a clean
133+
# generation pass, and only then a profiled pass of profiler_steps steps.
134+
enable_profiler: False
135+
profiler_steps: 5
136+
137+
# ML Diagnostics settings
138+
enable_ml_diagnostics: False
139+
profiler_gcs_path: ""
140+
enable_ondemand_xprof: False
141+
142+
# Directory for the Orbax cache of the converted pipeline (transformer + VAE +
143+
# text encoder). '' disables it. The first run loads from Diffusers and writes
144+
# the cache; later runs restore from it and skip the torch -> flax conversion.
145+
pretrained_orbax_dir: ""
146+
147+
# Keys below are required by pyconfig but unused for Z-Image inference.
148+
jax_cache_dir: ""
149+
dataset_name: ""
150+
dataset_save_location: "/mnt/disks/data/tmp"
151+
learning_rate_schedule_steps: 1
152+
max_train_steps: 1
153+
compile_topology_num_slices: -1
154+
quantization_local_shard_count: -1
155+
use_qwix_quantization: False

‎src/maxdiffusion/configs/base_zimage_turbo.yml‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# Official Z-Image-Turbo inference defaults. The base model uses the same
16+
# transformer; only its recommended denoising-step count differs.
17+
18+
run_name: 'zimage-turbo'
19+
20+
# If true save config to GCS in {base_output_directory}/{run_name}/
21+
save_config_to_gcs: False
22+
23+
pretrained_model_name_or_path: 'Tongyi-MAI/Z-Image-Turbo'
24+
model_name: z_image
25+
model_type: 'T2I'
26+
27+
unet_checkpoint: ''
28+
# This will convert the weights to this dtype.
29+
# When running inference on TPUv5e, use weights_dtype: 'bfloat16'
30+
weights_dtype: 'bfloat16'
31+
# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype)
32+
activations_dtype: 'bfloat16'
33+
34+
# Maximum sequence length for the text encoder
35+
max_sequence_length: 512
36+
37+
# Attention
38+
attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
39+
# Maxdiffusion has 2 types of attention sharding strategies:
40+
# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
41+
# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is
42+
# sharded in cross attention q.
43+
attention_sharding_uniform: True
44+
# Empty dict lets the kernel pick defaults. Example v6e override:
45+
# flash_block_sizes: {
46+
# "block_q" : 1024,
47+
# "block_kv_compute" : 1024,
48+
# "block_kv" : 1024,
49+
# "block_q_dkv" : 1024,
50+
# "block_kv_dkv" : 1024,
51+
# "block_kv_dkv_compute" : 1024,
52+
# "block_q_dq" : 1024,
53+
# "block_kv_dq" : 1024,
54+
# "use_fused_bwd_kernel": False,
55+
# }
56+
flash_block_sizes: {}
57+
58+
# Output directory
59+
# Create a GCS bucket, e.g. my-maxdiffusion-outputs and set this to "gs://my-maxdiffusion-outputs/"
60+
base_output_directory: ""
61+
output_dir: "/mnt/disks/data/tmp/maxdiffusion"
62+
# Local path for the generated image. If empty, an image named
63+
# {run_name}_output_{seed}.png is written to the working directory.
64+
output_file: "/mnt/disks/data/tmp/zimage-turbo.png"
65+
66+
# Hardware
67+
hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
68+
skip_jax_distributed_system: True
69+
70+
# Parallelism
71+
mesh_axes: ['data', 'fsdp', 'context', 'tensor']
72+
# Z-Image runs sequence parallel (activations sharded over `context`) with
73+
# data parallelism; the tensor axis is left at 1. `attention_sharding_uniform`
74+
# adds the matching q/kv sequence rules on top of these.
75+
logical_axis_rules: [
76+
['batch', ['data', 'fsdp']],
77+
['activation_batch', ['data', 'fsdp']],
78+
['activation_length', 'context'],
79+
['activation_heads', 'tensor'],
80+
['heads', 'tensor'],
81+
['mlp', 'tensor'],
82+
['embed', 'fsdp'],
83+
['norm', 'tensor'],
84+
['out_channels', 'tensor'],
85+
]
86+
data_sharding: [['data', 'fsdp', 'context', 'tensor']]
87+
88+
# One axis for each parallelism type may hold a placeholder (-1)
89+
# value to auto-shard based on available slices and devices.
90+
#
91+
# Measured on v6e-8, Turbo, 1024x1024, 9 steps, 8 images (per_device_batch_size 1.0):
92+
# data=8 2.45s 0.307s/image <- default
93+
# data=4 ctx=2 2.60s 0.325s/image
94+
# data=2 ctx=4 3.13s 0.391s/image
95+
# ctx=8 4.15s 0.519s/image
96+
# fsdp=8 9.43s 1.179s/image
97+
# The 12B of weights fit in HBM, so replicating them and giving each chip a
98+
# whole image beats every sharded layout: no collectives at all. fsdp re-gathers
99+
# every weight on all 9 denoising steps with no backward pass to amortize it.
100+
#
101+
# For single-image *latency* instead of throughput, generate one image
102+
# (per_device_batch_size: 0.125) with ici_data_parallelism: 1 and
103+
# ici_context_parallelism: -1, which splits one image's sequence 8 ways: 0.587s.
104+
dcn_data_parallelism: 1
105+
dcn_fsdp_parallelism: 1
106+
dcn_context_parallelism: 1
107+
dcn_tensor_parallelism: 1
108+
ici_data_parallelism: -1 # recommended ICI axis to be auto-sharded
109+
ici_fsdp_parallelism: 1
110+
ici_context_parallelism: 1
111+
ici_tensor_parallelism: 1
112+
113+
allow_split_physical_axes: False
114+
115+
# Generation parameters
116+
prompt: "A red fox reading a book in a quiet library, cinematic light"
117+
height: 1024
118+
width: 1024
119+
num_inference_steps: 9
120+
# The published Turbo configuration uses guidance_scale=0.
121+
guidance_scale: 0.0
122+
seed: 42
123+
# Images per device. Fractional values are allowed: on 8 devices,
124+
# 0.125 generates a single image.
125+
per_device_batch_size: 1.0
126+
# Latents decoded per VAE call. The decoder is replicated and its
127+
# activations are full-resolution, so raising this trades memory for speed.
128+
vae_decode_chunk: 1
129+
130+
# Profiling
131+
# generate_zimage runs an un-profiled warmup pass (compile), then a clean
132+
# generation pass, and only then a profiled pass of profiler_steps steps.
133+
enable_profiler: False
134+
profiler_steps: 5
135+
136+
# ML Diagnostics settings
137+
enable_ml_diagnostics: False
138+
profiler_gcs_path: ""
139+
enable_ondemand_xprof: False
140+
141+
# Directory for the Orbax cache of the converted pipeline (transformer + VAE +
142+
# text encoder). '' disables it. The first run loads from Diffusers and writes
143+
# the cache; later runs restore from it and skip the torch -> flax conversion.
144+
pretrained_orbax_dir: ""
145+
146+
# Keys below are required by pyconfig but unused for Z-Image inference.
147+
jax_cache_dir: ""
148+
dataset_name: ""
149+
dataset_save_location: "/mnt/disks/data/tmp"
150+
learning_rate_schedule_steps: 1
151+
max_train_steps: 1
152+
compile_topology_num_slices: -1
153+
quantization_local_shard_count: -1
154+
use_qwix_quantization: False

‎src/maxdiffusion/generate_flux2klein.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,8 @@
3737

3838
from maxdiffusion.models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
3939
from maxdiffusion.models.vae_flax import FlaxAutoencoderKL
40-
from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model, load_and_convert_qwen3_weights
40+
from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model
41+
from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
4142
from maxdiffusion.schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler
4243

4344

0 commit comments

Comments
 (0)