diff --git a/recipes/Qwen/Qwen-Image-2.1.md b/recipes/Qwen/Qwen-Image-2.1.md index 3d9295e60..19fd6b912 100644 --- a/recipes/Qwen/Qwen-Image-2.1.md +++ b/recipes/Qwen/Qwen-Image-2.1.md @@ -129,17 +129,23 @@ vllm serve Qwen/Qwen-Image-2.1 --omni --no-enable-cuda-graph-decode The same control is available as `Omni(..., enable_cuda_graph_decode=False)` and as `enable_cuda_graph_decode: false` on a stage in the deployment YAML. `--enforce-eager` / `enforce_eager: true` remains the broader switch: it -disables both graph capture and automatic `torch.compile`. Otherwise, this -transformer uses its model-specific decode graphs instead of automatic -`torch.compile`. +disables both graph capture and automatic `torch.compile`. The two +optimizations stack: the DiT blocks are regionally `torch.compile`'d and +graph capture records the compiled (fused) kernels, while inductor's own +cudagraphs stay disabled so the two graph layers never nest. Measured on a +single GB200 (1024×1024, 50 steps, seed 42, BF16, `true_cfg_scale=1.0`, +warmup 1 + median of 3): compile+graph 3.04 s, compile-only 3.2 s, +graph-only 4.0 s, eager 4.3 s end-to-end per image; combo output vs eager is +44.5–45.6 dB PSNR (the same magnitude as pure `torch.compile` fusion +divergence), and graph-only vs eager is bit-identical. The autoregressive engine's `compilation_config.cudagraph_mode` does not control this diffusion path; use `enable_cuda_graph_decode` (or `enforce_eager` for everything) to force eager decode execution. -TP/SP/ring parallelism, HSDP, offload/cache hooks, compiled blocks, quantized -KV caches, dynamic LoRA, padded text masks, and a second in-flight request -whose cache aliases an already-owned graph key fall back to eager decode. +TP/SP/ring parallelism, HSDP, offload/cache hooks, quantized KV caches, +dynamic LoRA, padded text masks, and a second in-flight request whose cache +aliases an already-owned graph key fall back to eager decode. Graphs keep separate entries for different image layouts and copy request-owned prefix K/V into static buffers before replay; each graph key has a single live owner at a time, so concurrent same-shape requests never @@ -374,7 +380,10 @@ filled in: - [ ] TP=2/4 output parity vs single-card — pending measurement. - [ ] VAE tiling vs non-tiled decode — pending measurement. -- [ ] CUDA graph decode vs eager — pending measurement. +- [x] CUDA graph decode vs eager — measured on GB200 (1024×1024, 50 steps, + seed 42, BF16): compile+graph 3.04 s/image vs eager 4.3 s/image (~1.4×), + 44.5–45.6 dB PSNR vs eager; graph-only is bit-identical to eager. See + "CUDA Graph decode" above. ## Known Limitations diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5281d7790..475185810 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -959,9 +959,10 @@ class OmniDiffusionConfig: # Currently only Qwen-Image-2.1's transformer implements this; other models # ignore the flag. This is a per-model diffusion path, independent of the # AR engine's ``compilation_config.cudagraph_mode`` (which does not apply - # to diffusion stages) and of ``diffusion_compile_granularity`` — models - # with decode graphs skip automatic torch.compile of their DiT blocks, so - # the two never stack. ``enforce_eager=True`` forces eager decode and + # to diffusion stages). It stacks with ``diffusion_compile_granularity``: + # the DiT blocks are still torch.compile'd and graph capture records the + # compiled (fused) kernels, while inductor's own cudagraphs stay off. + # ``enforce_eager=True`` forces eager decode and # disables capture regardless of this flag; unsupported configurations # (SP/TP, ring, HSDP, dynamic LoRA, padded masks, quantized prefix KV, or # a second in-flight request aliasing the same graph key) log and fall diff --git a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py index e3707516f..655e4d6cc 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py +++ b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py @@ -24,12 +24,16 @@ Memory note: graph entries own an additional copy of the prefix K/V. Decode attention concatenates that prefix with target K/V inside the captured region; transient tensors use the graph memory pool. +The model runner regionally torch.compile's the transformer blocks even when +decode graphs are enabled, so capture records the compiled (fused) kernels; +inductor's own cudagraphs stay off so the two graph layers never stack. + Fallbacks (all logged, all eager): CUDA unavailable, model-level CPU offload without persistent weight staging, sequence/ring/tensor parallelism, -HSDP/offload hooks, torch.compile'd blocks, KV-cache quantization, dynamic -LoRA wrappers, padded text masks (the masked attention path branches on -mask contents, which cannot be captured), unknown graph key at decode, a -second live request aliasing an owned graph key, and capture failure. +HSDP/offload hooks, KV-cache quantization, dynamic LoRA wrappers, padded text +masks (the masked attention path branches on mask contents, which cannot be +captured), unknown graph key at decode, a second live request aliasing an +owned graph key, and capture failure. """ from __future__ import annotations @@ -223,9 +227,6 @@ class QwenImage21DecodeGraphManager: reason = "tensor parallelism runs collectives inside the captured region" if reason is None: for block in model.transformer_blocks: - if getattr(block.forward, "_torchdynamo_orig_callable", None) is not None: - reason = "transformer blocks are torch.compile'd" - break if hasattr(block, "_hook_registry") or hasattr(block, "_omni_original_forward"): reason = "transformer blocks carry offload/cache hooks" break diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 49cb9452c..d9cc38481 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -240,8 +240,21 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): return if getattr(model, "enable_cuda_graph_decode", False): - logger.info("Model runner: %s uses CUDA graph decode instead of torch.compile.", attr_name) - return + # Decode-graph models still compile their blocks: graph capture + # records the compiled (fused) kernels. Inductor's own cudagraphs + # must stay off so the two graph layers never stack. + try: + import torch._inductor.config as inductor_config + + inductor_config.triton.cudagraphs = False + except Exception as e: + logger.warning( + "Model runner: could not disable inductor cudagraphs for the " + "compile+CUDA-graph combo (%s); capture may fail and fall back " + "to compiled eager decode.", + e, + ) + logger.info("Model runner: %s combines CUDA graph decode with torch.compile.", attr_name) compile_granularity = self.od_config.diffusion_compile_granularity compile_dynamic = self.od_config.diffusion_compile_dynamic