Regionally compile the DiT blocks even when enable_cuda_graph_decode is set
and let graph capture record the compiled (fused) kernels; inductor's own
cudagraphs stay off so the two graph layers never nest. Measured on GB200
(1024x1024, 50 steps, seed 42, BF16): combo 3.04s vs compile-only 3.2s,
graph-only 4.0s, eager 4.3s per image; 44.5-45.6 dB PSNR vs eager.
Match the Qwen-Image 2.0 split: decoder kwargs/attrs use the correct
temporal_upsample spelling while the checkpoint-facing temperal_downsample
config key is kept as-is. Fixes pre-commit typos rewriting the name.
- Normalize tuple image inputs in pre_process_func (empty tuple -> T2I path)
- Use torch.index_select in QwenImage21Rope for NPU compatibility (#6320 style)
- Refuse concurrent aliasing of shared decode-graph prefix KV entries: track a
single live owner per key, non-owners fall back to eager decode
* init
Signed-off-by: Isotr0py <mozf@inferact.ai>
* fix
Signed-off-by: Isotr0py <mozf@inferact.ai>
* test(diffusion): add Qwen-Image-2.1 unit tests
Signed-off-by: Isotr0py <mozf@inferact.ai>
* feat(qwen_image_21): opt-in CUDA graph capture for KV-cache decode steps
Decode steps (1..N) recompute only the target image's tokens at a fixed
shape, but the per-request dict KV cache made addresses unstable, so the
step could not be captured. Prefill now copies prefix K/V into persistent
fixed-address buffers (keyed by branch/batch/prefix_len/target_tokens/
dtype/device/backend) and re-points the caller's cache dicts at them;
decode replays a graph captured per key over the full 32-block forward.
Eager fallback (logged) for non-CUDA, SP/ring/TP>1, HSDP, compiled or
offload-hooked blocks, KV-cache quantization, dynamic LoRA wrappers,
padded text masks, unknown keys, and capture failure.
Off by default; opt in via OmniDiffusionConfig.enable_cuda_graph_decode
or the QwenImage21Transformer2DModel constructor kwarg.
E2E (GB200, T2I 1024x1024, 50 steps, CFG 4.0, CUDNN_ATTN):
171.2 -> 142.1 ms/step, 8.70s -> 7.24s total (1.20x), 44.9 dB PSNR
vs eager output.
* clean
Signed-off-by: Isotr0py <mozf@inferact.ai>
---------
Signed-off-by: Isotr0py <mozf@inferact.ai>
Isotr0py