From d3cf26c1b1cf942e23e07ecaed0840e2c090658a Mon Sep 17 00:00:00 2001
From: Isotr0py <Isotr0py@outlook.com>
Date: Mon, 21 Sep 2026 08:36:34 +0000
Subject: [PATCH] Wire enable_cuda_graph_decode through engine args, stage
 config and CLI

Follows the vae_use_tiling path: OmniDiffusionConfig field, OmniEngineArgs,
StageDeployConfig, _DiffusionConfigProjection and --[no-]enable-cuda-graph-decode.
Documents interplay with enforce_eager and cudagraph_mode.
---
 vllm_omni/config/omni_config.py    |  4 ++++
 vllm_omni/config/stage_config.py   |  4 ++++
 vllm_omni/diffusion/data.py        | 12 ++++++++++++
 vllm_omni/engine/arg_utils.py      |  4 ++++
 vllm_omni/entrypoints/cli/serve.py | 10 ++++++++++
 5 files changed, 34 insertions(+)

diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index d64372b84..2ef5c9a58 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -851,6 +851,10 @@ class _DiffusionConfigProjection:
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)
+    # Mirrors OmniDiffusionConfig.enable_cuda_graph_decode; enforced by the
+    # model (Qwen-Image-2.1) together with enforce_eager, which lives in the
+    # shared engine config.
+    enable_cuda_graph_decode: bool = True
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index aab4752e0..423c86bc8 100644
--- a/vllm_omni/config/stage_config.py
+++ b/vllm_omni/config/stage_config.py
@@ -462,6 +462,10 @@ class StageDeployConfig:
     # Diffusion execution, cache, and VAE behavior.
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    # CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1
+    # today). Independent of compilation_config.cudagraph_mode;
+    # enforce_eager=True also disables it.
+    enable_cuda_graph_decode: bool | None = None
     fa_deterministic: bool | None = None
     cache_backend: str | None = None
     cache_config: dict[str, Any] | None = None
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index c643f5ff8..5281d7790 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -955,6 +955,18 @@ class OmniDiffusionConfig:
 
     # Compilation
     enforce_eager: bool = False
+    # Capture fixed-shape KV-cache decode (denoising) steps into CUDA graphs.
+    # 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
+    # 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
+    # back to eager decode.
+    enable_cuda_graph_decode: bool = True
     # Controls the generic compilation path used when a pipeline does not
     # provide its own setup_compile() implementation.
     diffusion_compile_granularity: str = "regional"
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 939b19d50..65124322d 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -574,6 +574,10 @@ class OrchestratorArgs:
     diffusion_attention_config: str | None = None
     diffusion_compile_granularity: str | None = None
     diffusion_compile_dynamic: bool | None = None
+    # CUDA graph capture of fixed-shape KV-cache decode steps (Qwen-Image-2.1
+    # today). None defers to the OmniDiffusionConfig default (enabled);
+    # --enforce-eager also disables it.
+    enable_cuda_graph_decode: bool | None = None
     cache_backend: str = "none"
     cache_config: str | None = None
     video_output_transport: dict[str, object] | None = None
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 10ae2166d..e8ad89bf4 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -604,6 +604,16 @@ class OmniServeCommand(CLISubcommand):
                 "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
             ),
         )
+        omni_config_group.add_argument(
+            "--enable-cuda-graph-decode",
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help=(
+                "Capture fixed-shape KV-cache decode steps into CUDA graphs on models that implement them "
+                "(currently Qwen-Image-2.1). Independent of compilation_config.cudagraph_mode; "
+                "disable with --no-enable-cuda-graph-decode. --enforce-eager also disables it."
+            ),
+        )
         omni_config_group.add_argument(
             "--fa-deterministic",
             dest="fa_deterministic",