Skip to content

Z-Image on AMD GPUs: the 1344x768 bucket captured, batch sizes 1-8 - #1107

Draft
merceod wants to merge 10 commits into
users/atindra/pr-33-v4-fp8-ditfrom
users/atindra/pr-41-v7-zimage
Draft

merceod wants to merge 10 commits into
users/atindra/pr-33-v4-fp8-ditfrom
users/atindra/pr-41-v7-zimage

Conversation

@merceod

@merceod merceod commented Nov 1, 2026 •

Copy link
Copy Markdown
Collaborator

Stacked on #1099 (pr-33-v4-fp8-dit). Picks the commits of #1101 first; they drop out of this diff as those merge.

Why

z_image_turbo trailed vllm-omni on MI325X: in window J the tip served 1.068 / 1.149 images/s at c=1 / c=16, vllm-omni 1.090 / 1.149 in window H. Two things cost it. 1344x768, half the request set, wasn't captured and ran eager, 955 ms an image against 920 for captured 1024x1024. And c=16 steps average 5.7 rows, so every batch size up to 8 needs a graph.

What changed

Stacked on #1099 with #1101 picked first (it brings #1100).

  • The 1344x768 bucket captured and its VAE compiled.
  • Capture batch sizes 1 to 8, a batch padded with dummy rows up to the next captured size. Capturing a second size without the padding cut c=16 from 1.07 to 0.94.
  • Dynamo's recompile limit sized to the capture plan (128 for 64 graphs), for this forward alone. At the default 84 a size silently ran uncompiled as 2496 eager kernels, c=1 1.21 s an image. A test checks that the limit follows the capture plan.
  • tuned_linears's 54-size hipBLASLt table for the DiT GEMMs comes in and goes out again (eaefc6e, 4749ed2), so the net diff has no table; 4749ed2 removes dit_linear.json, the config line, its docs and its GPU test.
  • A config comment on the warm-start cost of 64 graphs.

Verification

  • Window J, 8x MI325X, exclusive (job 550), three timed runs in one session; vllm-omni from window H. images/s:
arm c=1 c=16 PSNR / floor
tip f96e97d 1.068 1.149 29.85 / 26.09 PASS
this PR's code (table off) 1.125 1.196 [1.155-1.237] 29.53 / 26.09 PASS
branch with the table 1.132 [1.051-1.136] 1.227 [1.206-1.234] 29.58 / 26.09 PASS
vllm-omni, window H 1.090 1.149 25.12-25.31 / 26.09 FAIL (shared node)
  • Table off over vllm-omni: 1.03x at c=1, 1.04x at c=16; over the tip +5.3 % and +4.1 %. The table bought under 1 % at c=1 and its c=16 range covers the table-off median, so it's out. The branch arm's c=1 p95 was 0.898 s against a 0.881 p50 (vllm-omni's p95 0.954): no stalls on the quiet node.
  • On the integrated tip d559763 (one GPU, one job at a time on an empty queue): 1.118 / 1.222 images/s against the previous tip's 1.060 / 1.126 (+5.5 % / +7.7 %), PSNR 29.53, and 40 of 40 images match J's table-off arm per image, so the rebase, the limit route and the table's removal give what J measured.
  • Correctness: per image against the tip, median +0.075 dB (range -4.21 to +8.22, the spread between two kernels' rounding). vllm-omni misses the floor by 0.6 to 1 dB in five runs, cause unverified (not its bf16 RoPE tables).
  • Comparable work: both run 8 Euler steps without CFG at the same sizes, padding and seeds, on AITER attention and bf16 hipBLASLt. Per image vllm-omni spends text 25.8 / step 99 / VAE 48.1 ms, this tree 20-22 / 96 / 54 (42 at 1344x768).
  • CPU Core failing set equal to the tip's (34) and GPU 110 passed, 2 skipped on the branch tree. On d559763: CPU Core 10880 tests, 34 failing, the same set; test/rocm 1965 passed, 9 skipped. Counts are from those trees, not this cut.

Not in this PR

  • Warm start goes from 75-96 s to 178-223 s for 64 graphs. A cold box adds the 1344x768 VAE autotune, over 10 minutes.
  • A fresh VAE autotune at 1024x1024 may save 12 ms (the seeded pick is 54 ms, 1344x768's 42).
  • FLUX.2 klein 4B and Cosmos3 t2i share the scaffold's capture and recompile changes: PSNR identical to the tip on an earlier tree, speed not comparable there and not rerun. Unmeasured on H100.
  • The 8 ms replay stalls seen on shared nodes come from neighbouring processes' host-memory evictions on this SR-IOV box, not from this code.

Rebased onto main

Rebased onto main (ff7819d, 2026-11-05); 6 of its 10 commits. Not here any more: z_image: rotate q/k in real arithmetic so Inductor fuses it with the norms (b7511c2) is on main as main's qk_norm_rope (#860, #972), the same goal by another route; z_image: text encoder can pad to the caption bucket instead of 512 (0512896) is dropped again by f087716 in this range (no served gain), so neither is here; 1762ce8, the tuned hipBLASLt table mechanism (opt-in per model) is removed again by 4749ed2 in this range (window J: the table off ties); only the table's own add and removal stay; z_image: drop the text padding option, no served gain to justify it (f087716) had nothing left to remove once 0512896 was skipped. The MI325X numbers here were measured on the AMD line before its rebase; the rebased line read level on the Qwen3-TTS, Qwen3-Omni, Qwen3.5 mtp2 and FLUX.2 klein smoke rows, and -8.7 % on the Qwen3.5-0.8B text c=64 row, which waits on the twin-slot re-port (phase 2).

amd-flashinfer 0.6.18+amd.1's generated AITER prefill config names __hip_bfloat16 before anything declares it, because the JIT's -DUSE_ROCM=1 turns off the hip_bf16.h include in torch's BFloat16.h. Every bf16 model on FlashInfer's auto backend dies at setup with a ninja log. Give the JIT the header through FLASHINFER_EXTRA_CUDAFLAGS on ROCm; 0.6.18+amd.2 includes it itself.
… pool's workers

The pool is built under a config patch at the top of warmup, since the setting is part of a graph's cache key and an env default would recompile everything a golden cache holds.
A submodule can compile at load, and klein does, so the pool built at the top of warmup came too late: the first compile had already built it under inductor's defaults, quiesce included. The call moves to the top of EngineManager.build, so no model has to ask for it.
…forward alone

Rebase note: the limit also scales with main's step schedule, whose bf16 twin compiles the same code objects.
Window J: table off ties it within 1 % at c=1 and inside the noise at c=16. The RoPE in real arithmetic, the 1344x768
capture, capture batch sizes 1-8 and the recompile limit stay.
@merceod merceod changed the title Z-Image on AMD GPUs: real-arithmetic RoPE, the 1344x768 bucket captured, batch sizes 1-8 Z-Image on AMD GPUs: the 1344x768 bucket captured, batch sizes 1-8 Nov 5, 2026
@merceod
merceod changed the base branch from users/atindra/pr-40-v5b-step-fp8 to users/atindra/pr-33-v4-fp8-dit November 5, 2026 17:54

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants