[None][feat] Enable PyTorch profiler traces for VisualGen - #16814
Open
chang-l wants to merge 12 commits into
Open
[None][feat] Enable PyTorch profiler traces for VisualGen#16814chang-l wants to merge 12 commits into
chang-l wants to merge 12 commits into
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughVisualGen profiling now supports environment-controlled PyTorch traces coordinated with CUDA profiler ranges, denoising-phase hooks, rank- and window-specific trace output, executor integration through ChangesVisualGen profiling and inference integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DiffusionExecutor
participant BasePipeline
participant VisualGenProfiler
participant CUDAProfiler
participant TorchProfiler
DiffusionExecutor->>BasePipeline: Call run_inference(req)
BasePipeline->>VisualGenProfiler: Enter request_scope()
VisualGenProfiler->>CUDAProfiler: Start selected window
VisualGenProfiler->>TorchProfiler: Start tracing
BasePipeline->>VisualGenProfiler: Iterate denoise steps
VisualGenProfiler->>CUDAProfiler: Stop profiling window
VisualGenProfiler->>TorchProfiler: Export Chrome trace
BasePipeline-->>DiffusionExecutor: Return PipelineOutput
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/pipeline.py (1)
191-212: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate profiler lifecycle methods.
Add
-> Noneto both methods.- def _cuda_profiler_start(self): + def _cuda_profiler_start(self) -> None: ... - def _cuda_profiler_stop(self): + def _cuda_profiler_stop(self) -> None:As per coding guidelines, “Annotate every function, use
Nonefor non-returning functions.”🤖 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 `@tensorrt_llm/_torch/visual_gen/pipeline.py` around lines 191 - 212, Add the `-> None` return annotation to both `_cuda_profiler_start` and `_cuda_profiler_stop`, preserving their existing profiler lifecycle behavior.Source: Coding guidelines
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 178-179: Update the trace path construction around
self._torch_profile_trace_path so each profiling window receives a unique output
filename instead of reusing the same rank-based path. Add and maintain a window
index across repeated A-B,C-D profiling stops, incorporating it into the
exported trace path while preserving the existing rank and extension components.
In `@tests/unittest/_torch/visual_gen/test_profiler.py`:
- Around line 17-88: The profiler tests cover only a single start/stop window;
add a test for multiple configured ranges such as 0-1 and 3-4, asserting the
corresponding profiling behavior across both windows. Also register
tests/unittest/_torch/visual_gen/test_profiler.py in the relevant test-db and QA
test-list files, preserving their existing list format.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 191-212: Add the `-> None` return annotation to both
`_cuda_profiler_start` and `_cuda_profiler_stop`, preserving their existing
profiler lifecycle behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ce070dc7-3ebc-4306-babd-3cbb37cc95e9
📒 Files selected for processing (3)
docs/source/developer-guide/perf-analysis.mdtensorrt_llm/_torch/visual_gen/pipeline.pytests/unittest/_torch/visual_gen/test_profiler.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/visual_gen/pipeline.py (2)
181-190: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winSelect the profiler activity at runtime.
torch.profiler.profile()should not be given bothCUDAandXPUunconditionally; on builds that only support one accelerator backend, profiler setup can fail. Build the list from the supported activities or branch on the active device so the trace requests only CPU + the matching accelerator.🤖 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 `@tensorrt_llm/_torch/visual_gen/pipeline.py` around lines 181 - 190, Update the profiler setup around self._torch_profiler and activities to select only the accelerator activity supported by the active device or build, always retaining CPU while excluding the incompatible CUDA/XPU activity. Pass the resulting CPU-plus-matching-accelerator list to torch.profiler.profile.
189-189: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
with_modules=Truewon’t record module hierarchy here — VisualGen uses eagernn.Modules, so the profiler trace won’t include module info unless this path is scripted/traced or you add explicit instrumentation.🤖 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 `@tensorrt_llm/_torch/visual_gen/pipeline.py` at line 189, The profiler configuration using with_modules=True does not capture hierarchy for VisualGen’s eager nn.Module execution. Remove this ineffective option or replace it with explicit module instrumentation, ensuring module hierarchy is recorded only through a supported mechanism.
🧹 Nitpick comments (1)
tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py (1)
226-253: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExpand profiler lifecycle coverage
test_forward_honors_profile_step_rangeonly covers the0-1path. Add focused cases forpredenoise,all,postdenoise, warmup gating, and the final inactive state; thepostdenoisecase should assert that tracing stops and exports.Coverage summary: added
test_forward_honors_profile_step_range; listed intests/integration/test_lists/test-db/l0_a10.yml; no matching entry intests/integration/test_lists/qa/. Coverage verdict: needs follow-up.🤖 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/unittest/_torch/visual_gen/test_qwen_image_pipeline.py` around lines 226 - 253, Expand test_forward_honors_profile_step_range with focused profiler lifecycle cases covering predenoise, all, postdenoise, warmup gating, and the final inactive state. Configure each profile range and assert the corresponding CUDA and torch profiler calls, especially that postdenoise stops tracing and exports the configured trace. Keep the tests isolated with the existing pipeline test doubles and mocks.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.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py`:
- Around line 528-530: Update the forward path around
_start_postdenoise_profile() so post-denoise profiling is stopped and exported
after _decode_latents() completes. Use a finally block to call the corresponding
stop operation even when decoding raises, ensuring profiler state is not left
active and the Chrome trace is exported.
- Around line 488-494: Move the `_start_predenoise_profile()` call earlier in
the request flow, before timestep creation and scheduler preparation, so
pre-loop scheduler refresh work is included in profiling. Keep
`_start_denoise_profile()` immediately before the denoising loop.
---
Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 181-190: Update the profiler setup around self._torch_profiler and
activities to select only the accelerator activity supported by the active
device or build, always retaining CPU while excluding the incompatible CUDA/XPU
activity. Pass the resulting CPU-plus-matching-accelerator list to
torch.profiler.profile.
- Line 189: The profiler configuration using with_modules=True does not capture
hierarchy for VisualGen’s eager nn.Module execution. Remove this ineffective
option or replace it with explicit module instrumentation, ensuring module
hierarchy is recorded only through a supported mechanism.
---
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py`:
- Around line 226-253: Expand test_forward_honors_profile_step_range with
focused profiler lifecycle cases covering predenoise, all, postdenoise, warmup
gating, and the final inactive state. Configure each profile range and assert
the corresponding CUDA and torch profiler calls, especially that postdenoise
stops tracing and exports the configured trace. Keep the tests isolated with the
existing pipeline test doubles and mocks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dbb8a4b9-421d-4cb2-bc91-1fad5661b0e9
📒 Files selected for processing (4)
docs/source/developer-guide/perf-analysis.mdtensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.pytensorrt_llm/_torch/visual_gen/pipeline.pytests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/source/developer-guide/perf-analysis.md
chang-l
requested review from
QiJune,
brnguyen2,
crazydemo,
jieli-matrix,
mlefeb01 and
mzweilz
July 24, 2026 06:35
chang-l
marked this pull request as draft
July 24, 2026 16:38
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Profiling state and lifecycle lived directly on BasePipeline: five attributes and six methods on a class that is already 1300+ lines. Every test double had to reproduce that state (the Qwen-Image doubles carried eight profiler attributes just to call forward()), and the profiler tests could only reach the logic by rebinding unbound methods onto SimpleNamespace fakes. Move it all to VisualGenProfiler in a new profiler.py. BasePipeline keeps one attribute and thin hooks that gate on warmup and delegate. Three changes ride along: * Close each capture window on an idle device. A collector may stop collecting -- or end the process, as nsys --capture-range-end=stop-shutdown does -- the moment the range closes, so torch.cuda.synchronize() must run first. PyExecutor.profile_step() already does this; VisualGen did not. * Name the hooks for what they do. _profile_denoise_start() stopped the profiler and _profile_denoise_end() started it, which reads backwards at the call site. They are now _close_predenoise_window() and _open_postdenoise_window(), matching _open_step_window() / _close_step_window(). * Test the profiler directly. No SimpleNamespace fakes, no MethodType rebinding; new coverage for the pre-stop sync, the no-op close, request scopes that unwind on an exception, and single-shot phases not re-arming. No behavior change beyond the added synchronize(). Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
visual_gen has three denoise loops: BasePipeline.denoise(), Qwen-Image's true-CFG loop, and LTX-2 two-stage's _refinement_denoise(). Only the first two were instrumented, so on LTX2TwoStagesPipeline a numeric TLLM_PROFILE_VISUAL_GEN_START_STOP range captured stage 1 and silently dropped stage 2. That is worse than the Qwen-Image gap it mirrors: Qwen produced no trace at all, which is obvious, whereas this produces a plausible trace missing half the denoising. Add the same hooks to the stage 2 loop, and add a source-level test that finds every loop in visual_gen stepping a transformer and asserts it calls them -- the next hand-written loop fails the test instead of quietly producing a short trace. Also document what multi-stage pipelines mean for the per-loop modes: numeric ranges now emit one trace file per stage, and postdenoise still opens at the end of stage 1, so its window spans the later stages too. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
A pipeline with its own denoise loop had to place four calls in the right
four spots -- close_predenoise before the loop, open/close_step_window
around each step, open_postdenoise after. Placing three of four still
compiles and still produces a trace, just a wrong one, which is the failure
mode that hid the Qwen-Image and LTX-2 stage 2 gaps.
Fold all four into the iterator the loop already needs:
for i, t in self._profile_denoise_steps(timesteps):
...
The generator closes the pre-denoise window before yielding the first step,
toggles step windows on the indices a numeric range names, and opens the
post-denoise window after the last one. Partial instrumentation is no
longer expressible, and the loop bodies keep their indentation. Breaking
out early or raising skips the post-denoise arm, as before -- request_scope
still closes any window left open.
All three denoise loops now differ from a plain enumerate() by one call,
and the structural test asserts exactly that.
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
steps() only closed a numeric window on the exact stop index, so a range
extending past a loop's last step never closed inside the loop --
request_scope() closed it at request exit instead. LTX-2 stage 2 runs three
steps, so 0-4 opened on its step 0 and stayed open through VAE decode:
stage2 step0 -> open -> step1 -> step2 -> [still open] -> decode -> close
The stage-2 trace therefore contained decode, contradicting the documented
one-trace-per-stage behavior, and on a pipeline with a third stage the
window would have swallowed that too.
Numeric ranges select denoise steps, so close any window still open when
the iterator exhausts. Regression test drives two stages with the second
shorter than the selected range. Keyword modes are untouched: the new close
is guarded on the range being numeric, so predenoise/postdenoise/all keep
their existing boundaries.
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
The synchronize() added when the profiler was extracted sat outside
close_window()'s try/finally. A sticky CUDA error from an earlier kernel
surfaces at exactly that call, and when it did:
- cudaProfilerStop() was never attempted, so an attached Nsight collector
kept its capture range open
- _active stayed True, so every later open_window() short-circuited and
every later close_window() re-raised -- profiling wedged for the rest
of the process
- the torch.profiler session stayed attached and recording
Verified all three against the previous commit before fixing.
Move the sync inside an outer try whose finally always runs the shutdown:
stop the torch profiler, export its trace, stop the CUDA gate, clear
_active. The CUDA error still propagates -- it should -- but it no longer
takes the profiler down with it.
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
chang-l
force-pushed
the
agent/visual-gen-torch-profiler
branch
from
July 27, 2026 18:25
310ef21 to
62f7864
Compare
chang-l
marked this pull request as ready for review
July 27, 2026 23:50
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/visual_gen/test_profiler.py (1)
1-437: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCoverage summary:
tests/integration/test_lists/test-db/l0_a10.ymlalready listsunittest/_torch/visual_gen/test_profiler.py. Verdict: insufficient — add asteps()test for a multi-pair range like0-1,3-4;test_each_window_uses_a_fresh_trace_fileonly covers repeatedopen_window()/close_window()calls.🤖 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/unittest/_torch/visual_gen/test_profiler.py` around lines 1 - 437, Add a focused test for VisualGenProfiler.steps using a multi-pair numeric range such as “0-1,3-4”, verifying events open and close at both configured windows while skipping the gap. Keep the existing test_step_windows_follow_numeric_ranges coverage unchanged and use _recording_profiler to assert the complete event sequence.Source: Path instructions
🧹 Nitpick comments (2)
tests/unittest/_torch/visual_gen/test_profiler.py (1)
135-169: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo symmetric test for
open_window()'s failure-cleanup path.
close_window()'s failure cleanup is covered twice (sync failure, export failure), butopen_window()'sexcept RuntimeErrorbranch (stoppingcudaProfilerStop()and clearing_torch_profilerwhentorch_profiler.start()/cudaProfilerStart()raises) has no test.Also applies to: 218-238
🤖 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/unittest/_torch/visual_gen/test_profiler.py` around lines 135 - 169, Add a unit test alongside test_close_window_stops_collectors_when_sync_fails covering open_window() failure cleanup: make torch_profiler.start() or cudaProfilerStart() raise RuntimeError, assert the error propagates, cudaProfilerStop() is called, profiler._torch_profiler is cleared, and profiler.active is false. Reuse the existing _profiler_with_torch_trace setup and mocking patterns.tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py (1)
221-254: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a true-CFG profiler case
true_cfg_scale=1.0with nonegative_promptstays on the non-CFG path, so this test doesn’t exercise Qwen-Image’s dual cond/uncond loop. Usetrue_cfg_scale > 1.0with a negative prompt, or add a second case, to cover the custom CFG path and its step window.🤖 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/unittest/_torch/visual_gen/test_qwen_image_pipeline.py` around lines 221 - 254, Update test_forward_honors_profile_step_range to exercise the true-CFG path by using true_cfg_scale greater than 1.0 and providing a negative_prompt, while preserving the existing profiler step-window assertions. Alternatively, add a separate test case covering the dual conditional/unconditional loop and its profiler start/stop behavior.
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/profiler.py`:
- Around line 165-197: Update the postdenoise handling in steps and its callers
so _postdenoise_pending is opened only after the final denoise loop, not after
each completed loop. Have BasePipeline.denoise and _refinement_denoise
communicate whether another denoise loop follows, preserving the existing
single-stage behavior while keeping Stage 2 denoising outside the
postdenoise/VAE-only window.
- Around line 203-260: Guard the CUDA runtime calls in open_window and
_shutdown_collectors with torch.cuda.is_available(), matching the existing guard
in close_window. Only call cudart().cudaProfilerStart() or cudaProfilerStop()
when CUDA is available, while preserving profiler startup, shutdown, and
state-reset behavior for XPU-only builds.
---
Outside diff comments:
In `@tests/unittest/_torch/visual_gen/test_profiler.py`:
- Around line 1-437: Add a focused test for VisualGenProfiler.steps using a
multi-pair numeric range such as “0-1,3-4”, verifying events open and close at
both configured windows while skipping the gap. Keep the existing
test_step_windows_follow_numeric_ranges coverage unchanged and use
_recording_profiler to assert the complete event sequence.
---
Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_profiler.py`:
- Around line 135-169: Add a unit test alongside
test_close_window_stops_collectors_when_sync_fails covering open_window()
failure cleanup: make torch_profiler.start() or cudaProfilerStart() raise
RuntimeError, assert the error propagates, cudaProfilerStop() is called,
profiler._torch_profiler is cleared, and profiler.active is false. Reuse the
existing _profiler_with_torch_trace setup and mocking patterns.
In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py`:
- Around line 221-254: Update test_forward_honors_profile_step_range to exercise
the true-CFG path by using true_cfg_scale greater than 1.0 and providing a
negative_prompt, while preserving the existing profiler step-window assertions.
Alternatively, add a separate test case covering the dual
conditional/unconditional loop and its profiler start/stop behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 6ac9098a-4d8b-4aba-8553-73c461096f4c
📒 Files selected for processing (10)
docs/source/developer-guide/perf-analysis.mdtensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.pytensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/profiler.pytests/integration/test_lists/test-db/l0_a10.ymltests/unittest/_torch/visual_gen/test_profiler.pytests/unittest/_torch/visual_gen/test_qwen_image_pipeline.pytests/unittest/_torch/visual_gen/test_visual_gen_params.py
🚧 Files skipped from review as they are similar to previous changes (4)
- docs/source/developer-guide/perf-analysis.md
- tensorrt_llm/_torch/visual_gen/executor.py
- tests/unittest/_torch/visual_gen/test_visual_gen_params.py
- tests/integration/test_lists/test-db/l0_a10.yml
Comment on lines
+165
to
+197
| def steps(self, timesteps: Iterable[Any]) -> Iterator[Tuple[int, Any]]: | ||
| """Enumerate one denoise loop's steps, driving every window it owns. | ||
|
|
||
| Wrapping the iterator rather than exposing separate | ||
| before-loop/per-step/after-loop hooks means a pipeline cannot | ||
| instrument half a loop: the phase boundaries ride along with the | ||
| enumeration a denoise loop already needs. | ||
| """ | ||
| if self.range == "predenoise": | ||
| self.close_window() | ||
|
|
||
| starts, stops = self.range if isinstance(self.range, tuple) else (frozenset(), frozenset()) | ||
| for i, t in enumerate(timesteps): | ||
| if i in starts: | ||
| self.open_window() | ||
| yield i, t | ||
| if i in stops: | ||
| self.close_window() | ||
|
|
||
| # Everything below is reached only when the loop runs to completion. | ||
| # A loop that raised or broke out early has no post-denoise work | ||
| # worth capturing, and ``request_scope`` closes whatever window is | ||
| # still open. | ||
| if starts: | ||
| # A numeric range selects denoise steps, so it must not outlive | ||
| # the loop. Without this, a range whose stop index is past the | ||
| # last step -- 0-4 against LTX-2 stage 2's three steps -- would | ||
| # stay open through VAE decode and, on a multi-stage pipeline, | ||
| # swallow the following stage as well. | ||
| self.close_window() | ||
| if self._postdenoise_pending: | ||
| self.open_window() | ||
| self._postdenoise_pending = False |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
postdenoise opens after the first denoise loop, not the last, on multi-stage pipelines.
_postdenoise_pending fires as soon as any steps() call runs to completion (Line 195-197). For a single-stage pipeline that's correct, but LTX2TwoStagesPipeline runs Stage 1 through BasePipeline.denoise() (which uses this same steps()) and then Stage 2 through its own loop in pipeline_ltx2_two_stages.py (also wired to _profile_denoise_steps). With postdenoise mode, the single-shot window opens right after Stage 1 finishes — swallowing all of Stage 2's refinement denoising into what's documented as a "VAE decode only" window (line 55-56: "profile from the end of the last denoise loop to request completion, covering VAE decode"). Contrast with the numeric-range case, which the docstring explicitly calls out as needing per-stage handling (lines 70-74), but postdenoise has no equivalent multi-stage guard.
🐛 Possible direction
- if self._postdenoise_pending:
- self.open_window()
- self._postdenoise_pending = False
+ if self._postdenoise_pending and is_final_loop:
+ self.open_window()
+ self._postdenoise_pending = FalseThis would need callers (BasePipeline.denoise() for stage 1, _refinement_denoise() for stage 2) to tell steps() whether more denoise loops follow the current one.
🤖 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 `@tensorrt_llm/_torch/visual_gen/profiler.py` around lines 165 - 197, Update
the postdenoise handling in steps and its callers so _postdenoise_pending is
opened only after the final denoise loop, not after each completed loop. Have
BasePipeline.denoise and _refinement_denoise communicate whether another denoise
loop follows, preserving the existing single-stage behavior while keeping Stage
2 denoising outside the postdenoise/VAE-only window.
jieli-matrix
requested review from
yingguo-trt
and removed request for
jieli-matrix
July 28, 2026 01:52
yingguo-trt
approved these changes
Jul 28, 2026
The parse_profile_range docstring claimed postdenoise ran "from the end of the last denoise loop", but it is single-shot and arms at the end of the *first* one. On LTX-2 two-stage that window is a superset of VAE decode -- it also holds the spatial upsample and all of stage 2 -- which contradicted the docstring while the perf-analysis guide described it correctly. Keep the behavior and state the limitation in both places. It is a superset, not a gap: decode is always in the trace, just not alone. No mode isolates decode on that pipeline; doing so needs the profiler to know which loop is last, which the base denoise() cannot tell it. Add a test pinning the window contents so the behavior cannot drift without the docs being revisited. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
…rch-profiler Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/visual_gen/pipeline.py
Merging main brought in Qwen-Image-Edit (NVIDIA#16095) and Qwen-Image-Layered, each with its own denoise loop that bypasses BasePipeline.denoise(). Both were invisible to TLLM_PROFILE_VISUAL_GEN_START_STOP, the same gap this branch fixed for Qwen-Image and LTX-2 stage 2. test_every_denoise_loop_is_instrumented caught them on the merge, which is what it was written for. Route both through _profile_denoise_steps(). Also relax that test's exact loop count to a lower bound: a hard count fails on every newly added pipeline rather than on the thing that matters, which is whether each detected loop is instrumented. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
tensorrt_llm/_torch/visual_gen/profiler.py (1)
214-220: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard CUDA runtime calls on non-CUDA builds.
XPU tracing is explicitly supported in
_create_torch_profiler(), but these unconditionaltorch.cuda.cudart()calls fail before startup or during cleanup on XPU-only builds. Gate both start/stop calls withtorch.cuda.is_available()while still starting/stopping the torch profiler.Also applies to: 260-263
🤖 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 `@tensorrt_llm/_torch/visual_gen/profiler.py` around lines 214 - 220, Update the profiler start and cleanup logic around the CUDA runtime calls to check torch.cuda.is_available() before invoking torch.cuda.cudart(), cudaProfilerStart(), or cudaProfilerStop(). Keep _torch_profiler.start() and its corresponding stop behavior active regardless of CUDA availability so XPU-only tracing continues to work.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/profiler.py (1)
27-38: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse parameterized Python 3.10+ annotations.
ProfileRangeleaves its endpoint types unknown and uses legacyUnion/Tuple/Optionalforms. Usestr | tuple[frozenset[int], frozenset[int]] | Noneandstr | None.Also applies to: 133-133
🤖 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 `@tensorrt_llm/_torch/visual_gen/profiler.py` around lines 27 - 38, Update the ProfileRange alias to use Python 3.10 union and built-in generic syntax, explicitly typing both range endpoints as frozenset[int] and allowing None. Also replace the separately referenced Optional annotation at the indicated usage with str | None, removing legacy Union, Tuple, and Optional imports if no longer needed.Sources: Coding guidelines, Learnings
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/profiler.py`:
- Around line 49-59: Replace every en dash character in the profiler mode
docstring entries with Ruff-compliant punctuation, such as a hyphen or
equivalent ASCII separator, while preserving the documented meanings and
formatting of the A-B, range, predenoise, postdenoise, all, and unset options.
---
Duplicate comments:
In `@tensorrt_llm/_torch/visual_gen/profiler.py`:
- Around line 214-220: Update the profiler start and cleanup logic around the
CUDA runtime calls to check torch.cuda.is_available() before invoking
torch.cuda.cudart(), cudaProfilerStart(), or cudaProfilerStop(). Keep
_torch_profiler.start() and its corresponding stop behavior active regardless of
CUDA availability so XPU-only tracing continues to work.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/profiler.py`:
- Around line 27-38: Update the ProfileRange alias to use Python 3.10 union and
built-in generic syntax, explicitly typing both range endpoints as
frozenset[int] and allowing None. Also replace the separately referenced
Optional annotation at the indicated usage with str | None, removing legacy
Union, Tuple, and Optional imports if no longer needed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: da9a2e31-3c51-4289-972c-59c6e8e9c49d
📒 Files selected for processing (7)
docs/source/developer-guide/perf-analysis.mdtensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.pytensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.pytensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/profiler.pytests/integration/test_lists/test-db/l0_a10.yml
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/source/developer-guide/perf-analysis.md
fredricz-20070104
approved these changes
Jul 30, 2026
fredricz-20070104
left a comment
fredricz-20070104
left a comment
Collaborator
There was a problem hiding this comment.
Review summary - Approve
Reviewed the full diff; no blocking or major issues found.
Minor, non-blocking notes:
tensorrt_llm/_torch/visual_gen/profiler.py: CUDA runtime gate calls are unconditional while torch profiler supports XPUtensorrt_llm/_torch/visual_gen/profiler.py: Legacy typing forms / missing parameterization
Automated review by NVCortex Lite, run by @fredricz-20070104.
zhenhuaw-me
approved these changes
Jul 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Dev Engineer Review
TLLM_TORCH_PROFILE_TRACEis set, gated/scoped byTLLM_PROFILE_VISUAL_GEN_START_STOP.all,predenoise,postdenoise, and numeric denoise-step range selections.record_shapes.postdenoiseis positioned to cover the correct denoise boundary for pipelines with multiple stages (may extend beyond VAE decode when later stages run).VisualGenProfiler(tensorrt_llm/_torch/visual_gen/profiler.py) and refactored VisualGen pipeline profiling to use it.BasePipeline.run_inference()now wraps non-warmup requests inrequest_scope()and centralizes denoise-step windowing viaBasePipeline._profile_denoise_steps(...).tensorrt_llm/_torch/visual_gen/pipeline.py.self._profile_denoise_steps(timesteps)/ step-range iterators.pipeline_qwen_image.py,pipeline_qwen_image_edit.py,pipeline_qwen_image_layered.py).pipeline_ltx2_two_stages.py) so Stage 2 is the one driving the profiling capture window(s).pipeline.run_inference()rather thanpipeline.infer().docs/source/developer-guide/perf-analysis.md) with VisualGen-specific instructions, including howTLLM_PROFILE_VISUAL_GEN_START_STOPmaps to phases/ranges and trace export behavior (per-rank/per-window suffixing).QA Engineer Review
Test files modified
tests/unittest/_torch/visual_gen/test_profiler.pytests/unittest/_torch/visual_gen/test_qwen_image_pipeline.pytests/unittest/_torch/visual_gen/test_visual_gen_params.pytests/unittest/visual_gen/test_executor_shared_tensor_ipc.pytests/integration/test_lists/test-db/l0_a10.ymlTest functions added/updated
tests/unittest/_torch/visual_gen/test_profiler.pytest_setup_torch_profilertest_setup_torch_profiler_requires_profile_rangetest_window_controls_torch_profilertest_close_window_syncs_before_ending_capturetest_close_window_stops_collectors_when_sync_failstest_close_window_is_a_noop_when_never_openedtest_each_window_uses_a_fresh_trace_filetest_cuda_gate_closes_when_trace_export_failstest_request_scope_owns_phase_boundaries(parameterized)test_request_scope_closes_window_when_inference_raisestest_single_shot_phases_do_not_rearmtest_step_windows_follow_numeric_rangestest_numeric_window_does_not_outlive_a_short_denoise_looptest_run_inference_skips_profiling_during_warmuptest_every_denoise_loop_is_instrumented(AST structural guard for denoise loops)l0_a10.yml(see below) and should validate trace export/gating and lifecycle semantics.tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.pytest_forward_honors_profile_step_rangetests/unittest/_torch/visual_gen/test_visual_gen_params.pypipeline.run_inference().tests/unittest/visual_gen/test_executor_shared_tensor_ipc.pyrun_inference(req)delegating to existinginfer(req).Test-list coverage
tests/integration/test_lists/test-db/l0_a10.yml# visual_gen, includesunittest/_torch/visual_gen/test_profiler.py(added entry per PR change summary).Verdict: sufficient — VisualGen profiler unit tests are registered in
l0_a10.yml, so the added profiling lifecycle and trace-export behavior should be exercised in CI.++ BrianLi23 for viz
Description
VisualGen already uses
TLLM_PROFILE_VISUAL_GEN_START_STOPto gate Nsight Systems capture, but it could not emit PyTorch/Kineto traces like the LLM executor.This change starts an optional
torch.profiler.profilesession on the same VisualGen ranges whenTLLM_TORCH_PROFILE_TRACEis set:allnow covers the complete user request from text encoding through denoising and VAE decode.predenoisecovers request start through the denoise-loop boundary;postdenoisecovers that boundary through request completion.BasePipeline.run_inference()wrapper owns request-level start/stop and exception cleanup. Base denoise and Qwen-Image's custom true-CFG loop call the same semantic phase/step hooks, so Qwen does not duplicate profiler setup or teardown.cudaProfilerStart()/cudaProfilerStop()remain capture gates for an attached Nsight collector; they do not start Nsight on their own. Torch and Nsight collectors should be run in separate invocations because both use CUPTI.Test Coverage
test_profiler.pyinl0_a10.yml.13 passed, 38 warnings in 143.25s.error=NoneusingTLLM_PROFILE_VISUAL_GEN_START_STOP=all.PR Checklist
PR description clearly explains what and why.
PR follows the TRT-LLM coding guidelines.
Test cases are provided for the new code paths and are routed in the test database.
No public API changes or new dependencies.
Documentation is updated.
No CODEOWNERS or architecture-diagram change is required.
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, comment
/bot help.