[None][fix] Fix Qwen image CUDA graph with CFG (NVIDIA#16391)

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>
Signed-off-by: Yibin-Li <109242046+yibinl-nvidia@users.noreply.github.com>

Author: yibinl-nvidia

tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py

@@ -486,6 +486,8 @@ def forward(
         # Denoise loop.
         timer.mark_denoise_start()
         logger.info("Denoising (%d steps)...", len(timesteps))
+        pipeline_config = getattr(self, "pipeline_config", None)
+        cuda_graph_enabled = getattr(getattr(pipeline_config, "cuda_graph", None), "enable", False)
         for i, t in enumerate(timesteps):
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
             noise_pred = self.transformer(
@@ -497,6 +499,11 @@ def forward(
                 return_dict=False,
             )[0]
 
+            if do_true_cfg and cuda_graph_enabled:
+                # CUDA graph outputs are graph-owned buffers; the negative CFG
+                # replay may reuse the same pool before guidance consumes this one.
+                noise_pred = noise_pred.clone()
+
             if do_true_cfg:
                 neg_noise_pred = self.transformer(
                     hidden_states=latents,

tests/integration/defs/examples/visual_gen/test_visual_gen.py

@@ -811,17 +811,18 @@ def wan22_bf16_video_path(_visual_gen_deps, llm_venv):
     return output_path
 
 
-def _generate_qwenimage_lpips_image(model_path, output_path):
-    """Generate the QwenImage text-to-image LPIPS sample (default setting, compile-off)."""
+def _generate_qwenimage_lpips_image(model_path, output_path, *, enable_cuda_graph=False):
+    """Generate the QwenImage text-to-image LPIPS sample."""
     from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader
     from tensorrt_llm.media.encoding import save_image
-    from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs
+    from tensorrt_llm.visual_gen.args import CudaGraphConfig, TorchCompileConfig, VisualGenArgs
 
     _skip_if_missing(model_path, "QwenImage checkpoint", is_dir=True)
     _disable_inductor_compile_worker_quiesce()
     args = VisualGenArgs(
         model=model_path,
         torch_compile_config=TorchCompileConfig(enable=False),
+        cuda_graph_config=CudaGraphConfig(enable=enable_cuda_graph),
     )
     pipeline = PipelineLoader(args).load(skip_warmup=True)
     try:
@@ -1082,6 +1083,28 @@ def test_qwenimage_lpips_against_golden(tmp_path):
     _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_qwenimage_cuda_graph_lpips_against_golden(tmp_path):
+    generated_path = tmp_path / "qwenimage_cuda_graph_generated.png"
+    golden_path = _golden_media_path(
+        tmp_path, "qwenimage_lpips_golden.png", "QwenImage LPIPS golden image"
+    )
+    _generate_qwenimage_lpips_image(
+        _lpips_model_path(QWENIMAGE_MODEL_SUBPATH),
+        generated_path,
+        enable_cuda_graph=True,
+    )
+    score = _run_lpips_eval(
+        tmp_path,
+        "qwenimage_cuda_graph",
+        "image",
+        QWENIMAGE_LPIPS_PROMPT,
+        golden_path,
+        generated_path,
+    )
+    _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
 def test_cosmos3_nano_t2v_lpips_against_golden(_visual_gen_deps, tmp_path):
     generated_path = tmp_path / "cosmos3_nano_t2v_generated.mp4"

tests/integration/test_lists/test-db/l0_b200.yml

@@ -367,6 +367,7 @@ l0_b200:
   # Measured on B200 with TRT-LLM commit 85665f5f from the staging main image:
   # QwenImage ~5 min, Cosmos3-Nano T2I ~3 min, and T2V ~5 min.
   - examples/visual_gen/test_visual_gen.py::test_qwenimage_lpips_against_golden TIMEOUT (10)
+  - examples/visual_gen/test_visual_gen.py::test_qwenimage_cuda_graph_lpips_against_golden TIMEOUT (10)
   - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2i_lpips_against_golden TIMEOUT (10)
   - examples/visual_gen/test_visual_gen.py::test_cosmos3_nano_t2v_lpips_against_golden TIMEOUT (15)
   - visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark