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>

vllm_omni/diffusion/data.py

@@ -902,6 +902,11 @@ class OmniDiffusionConfig:
 
     # Compilation
     enforce_eager: bool = False
+    # Opt-in: capture fixed-shape KV-cache decode (denoising) steps into CUDA
+    # graphs. Currently implemented by Qwen-Image-2.1's transformer; unsupported
+    # configurations (SP/TP, ring, HSDP, dynamic LoRA, padded masks) log and
+    # fall back to eager decode. Other models ignore this flag.
+    enable_cuda_graph_decode: bool = False
     # Controls the generic compilation path used when a pipeline does not
     # provide its own setup_compile() implementation.
     diffusion_compile_granularity: str = "regional"