[None][fix] Fix Qwen image CUDA graph with CFG

[yibinl-nvidia]

Summary by CodeRabbit

• Bug Fixes

  • Improved QwenImage generation reliability when CUDA Graphs are enabled, preventing potential output-buffer conflicts during guided generation.

• Tests

  • Added integration coverage to verify CUDA Graph image generation quality against the expected reference output.

Description

Fix CUDA graph generation for Qwen Image when true CFG is enabled with a negative prompt. The images below compare eager mode and CUDA graph output after the fix.

Eager mode	Cuda graph on

CUDA graph vs eager LPIPS is 0.000000.

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #59289 [ run ] triggered by Bot. Commit: 0873a0f Link to invocation

[tensorrt-cicd]

PR_Github #59289 [ run ] completed with state FAILURE. Commit: 0873a0f
/LLM/main/L0_MergeRequest_PR pipeline #47775 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

QwenImage CUDA graph denoising now clones CFG predictions before buffer reuse, and integration coverage generates a CUDA-graph-enabled image and compares it with the existing LPIPS golden-image threshold.

Changes

QwenImage CUDA graph support

Layer / File(s)	Summary
Preserve CFG predictions during graph replay tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py	The pipeline reads CUDA graph configuration and clones noise_pred before negative CFG replay when true CFG is active.
Validate CUDA graph image generation tests/integration/defs/examples/visual_gen/test_visual_gen.py	The QwenImage LPIPS helper accepts CUDA graph configuration, and a new test evaluates CUDA-graph output against the existing golden-image threshold.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: karljang, bowenfu, niukuo, mzweilz, chang-l

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the change set and follows the required [None][fix] template.
Description check	✅ Passed	The description explains the bug and fix and broadly follows the template, though the Test Coverage section is left empty.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[yibinl-nvidia]

/bot run

[coderabbitai]

🧹 Nitpick comments (2)

tests/integration/defs/examples/visual_gen/test_visual_gen.py (2)

814-815: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate the changed functions.

Both _generate_qwenimage_lpips_image and test_qwenimage_cuda_graph_lpips_against_golden lack parameter and return annotations.

Proposed annotations

-def _generate_qwenimage_lpips_image(model_path, output_path, *, enable_cuda_graph=False):
+def _generate_qwenimage_lpips_image(
+    model_path: str | Path,
+    output_path: str | Path,
+    *,
+    enable_cuda_graph: bool = False,
+) -> None:
...
-def test_qwenimage_cuda_graph_lpips_against_golden(tmp_path):
+def test_qwenimage_cuda_graph_lpips_against_golden(tmp_path: Path) -> None:

As per coding guidelines, “Annotate every function.”

Also applies to: 1086-1087

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py` around lines
814 - 815, Annotate both _generate_qwenimage_lpips_image and
test_qwenimage_cuda_graph_lpips_against_golden with parameter and return type
annotations, covering every declared argument and the function return value. Use
the types consistent with surrounding test helpers and preserve the existing
behavior and signatures otherwise.

Source: Coding guidelines

---

1086-1105: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy lift

Add a direct eager-vs-CUDA-graph comparison.

This test only compares the graph output with qwenimage_lpips_golden.png using the existing 0.05 threshold. It does not directly verify that CUDA-graph output matches eager output, which is the behavior this fix targets; a regression could still pass the indirect golden-image check.

Generate an eager reference in this test and run LPIPS against the CUDA-graph result, preferably with a dedicated tighter threshold.

Suggested coverage extension

 def test_qwenimage_cuda_graph_lpips_against_golden(tmp_path):
+    eager_path = tmp_path / "qwenimage_eager_generated.png"
+    _generate_qwenimage_lpips_image(
+        _lpips_model_path(QWENIMAGE_MODEL_SUBPATH),
+        eager_path,
+        enable_cuda_graph=False,
+    )
+
     generated_path = tmp_path / "qwenimage_cuda_graph_generated.png"
...
     _assert_lpips_below_threshold(score, QWENIMAGE_LPIPS_THRESHOLD)
+
+    eager_score = _run_lpips_eval(
+        tmp_path,
+        "qwenimage_cuda_graph_vs_eager",
+        "image",
+        QWENIMAGE_LPIPS_PROMPT,
+        eager_path,
+        generated_path,
+    )
+    _assert_lpips_below_threshold(eager_score, QWENIMAGE_LPIPS_THRESHOLD)

As per path instructions, this integration-test coverage is insufficient for the stated eager-vs-CUDA-graph objective and needs a concrete follow-up in tests/integration/defs/examples/visual_gen/test_visual_gen.py.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py` around lines
1086 - 1105, Extend test_qwenimage_cuda_graph_lpips_against_golden to generate
an eager reference image alongside the CUDA-graph output, then run LPIPS
directly between those two generated images. Assert the comparison with a
dedicated, tighter eager-vs-CUDA-graph threshold while preserving the existing
golden-image coverage as appropriate.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/integration/defs/examples/visual_gen/test_visual_gen.py`:
- Around line 814-815: Annotate both _generate_qwenimage_lpips_image and
test_qwenimage_cuda_graph_lpips_against_golden with parameter and return type
annotations, covering every declared argument and the function return value. Use
the types consistent with surrounding test helpers and preserve the existing
behavior and signatures otherwise.
- Around line 1086-1105: Extend test_qwenimage_cuda_graph_lpips_against_golden
to generate an eager reference image alongside the CUDA-graph output, then run
LPIPS directly between those two generated images. Assert the comparison with a
dedicated, tighter eager-vs-CUDA-graph threshold while preserving the existing
golden-image coverage as appropriate.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d75e6b9a-37c0-4445-9fb0-0596920df915

📥 Commits

Reviewing files that changed from the base of the PR and between 58d8964 and 7c24335.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
• tests/integration/defs/examples/visual_gen/test_visual_gen.py

[tensorrt-cicd]

PR_Github #59295 [ run ] triggered by Bot. Commit: 7c24335 Link to invocation

[yibinl-nvidia]

/bot kill

[tensorrt-cicd]

PR_Github #59297 [ kill ] triggered by Bot. Commit: 9f986e5 Link to invocation

[tensorrt-cicd]

PR_Github/16391-7c24335 #59295 was force-killed by a newer pipeline run.
L0 job information not available (job may not have been triggered yet).

Link to superseding invocation

[tensorrt-cicd]

PR_Github #59297 [ kill ] completed with state SUCCESS. Commit: 9f986e5
Successfully killed previous jobs for commit 9f986e5

Link to invocation

[yibinl-nvidia]

/bot run

[tensorrt-cicd]

PR_Github #59299 [ run ] triggered by Bot. Commit: 9f986e5 Link to invocation

[fredricz-20070104]

Approved.

[tensorrt-cicd]

PR_Github #59299 [ run ] completed with state SUCCESS. Commit: 9f986e5
/LLM/main/L0_MergeRequest_PR pipeline #47784 completed with status: 'SUCCESS'

CI Report

Link to invocation

[yibinl-nvidia]

/bot run