Files
vllm-omni/vllm_omni/diffusion
Isotr0py 4fc3cf8c60 Enable CUDA Graph for DiT decode layers (#4)
* 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>
2026-09-18 10:02:39 +08:00
..