[None][feat] Enable PyTorch profiler traces for VisualGen - #16814
Merged
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>
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
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62839 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62839 [ run ] completed with state |
…rch-profiler Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py # tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62882 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62882 [ run ] completed with state |
arysef
approved these changes
Jul 31, 2026
luyiyun1021
approved these changes
Jul 31, 2026
venkywonka
approved these changes
Jul 31, 2026
…rch-profiler Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/visual_gen/executor.py
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63074 [ run ] triggered by Bot. Commit: |
chang-l
enabled auto-merge (squash)
July 31, 2026 17:12
Collaborator
|
PR_Github #63074 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #63129 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63129 [ run ] completed with state |
yuanjingx87
pushed a commit
that referenced
this pull request
Aug 1, 2026
Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
1 task
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
8 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.
++ BrianLi23 for viz
Description
VisualGen already used
TLLM_PROFILE_VISUAL_GEN_START_STOPto gate Nsight Systems capture, but it could not emit PyTorch/Kineto traces the way the LLM executor does. This adds that, and fixes the coverage gaps found while wiring it up.PyTorch traces on the existing VisualGen ranges
When
TLLM_TORCH_PROFILE_TRACEis set, atorch.profiler.profilesession opens and closes on the same windows the CUDA/Nsight gate already uses:all— the complete request, text encoding through VAE decodepredenoise— request start through denoise-loop setuppostdenoise— end of a denoise loop through request completionA-B,A-B,C-D,A,B— per-denoise-loop step rangesWarmup is excluded. Each rank writes its own file (
-rank-N), and each repeated window within a rank gets a-window-Nsuffix, so nothing is overwritten. There are no public API or dependency changes.cudaProfilerStart()/cudaProfilerStop()remain capture gates for an attached Nsight collector — they do not start Nsight on their own. Run Nsight and torch traces in separate invocations, since both use CUPTI.Profiler ownership
VisualGenProfiler(_torch/visual_gen/profiler.py) owns every window decision.BasePipelineholds one instance and exposes two thin seams:run_inference()wraps a request inrequest_scope(), which opens theall/predenoisewindow and guarantees every window closes on the way out, including after an exception.DiffusionExecutorcalls this instead ofinfer()._profile_denoise_steps()replaces a denoise loop'senumerate()and drives every boundary that loop owns — closing the pre-denoise window before the first step, toggling step windows on the configured indices, and arming the post-denoise window after the last step.Folding all boundaries into the iterator means a loop cannot be half-instrumented. Each denoise loop differs from a plain
enumerate()by exactly one call.Denoise loops that bypass
BasePipeline.denoise()Five loops exist in
visual_gen; four do not go through the basedenoise()and were invisible to the profiler. All are now instrumented:BasePipeline.denoise()pipeline_qwen_image.py::forwardpipeline_ltx2_two_stages.py::_refinement_denoisepipeline_qwen_image_edit.py::forwardpipeline_qwen_image_layered.py::forwardtest_every_denoise_loop_is_instrumentedwalks thevisual_genAST, finds every loop that steps a transformer, and asserts each one uses_profile_denoise_steps(). It caught the two Qwen-Image loops above when main was merged in, which is what it was written for.Window lifecycle fixes
torch.cuda.synchronize()runs beforecudaProfilerStop(). A collector may stop collecting — or end the process, asnsys --capture-range-end=stop-shutdowndoes — the moment the range closes, so no async work may still be in flight. MirrorsPyExecutor.profile_step().0-4against LTX-2 stage 2's three steps) stayed open through VAE decode and, on a multi-stage pipeline, into the following stage. It now closes when the iterator exhausts.synchronize()previously skippedcudaProfilerStop(), left_activeset, and left Kineto recording — wedging the profiler for the rest of the process. The shutdown now runs from afinally; the error still propagates.Known limitation
postdenoiseis single-shot and arms after the first denoise loop. On a multi-stage pipeline its window is therefore a superset of VAE decode — on LTX-2 two-stage it also holds the spatial upsample and all of stage 2. Isolating decode there would require the profiler to know which loop is last, whichBasePipeline.denoise()cannot tell it. Documented inparse_profile_rangeand the perf-analysis guide, and pinned bytest_postdenoise_arms_after_the_first_loop_on_multi_stageso it cannot drift silently.Test Coverage
tests/unittest/_torch/visual_gen/test_profiler.py(registered inl0_a10.yml) covers:cudaProfilerStop(), and collector shutdown when the sync raisesall/predenoise/postdenoiseboundary ordering, single-shot phases not re-arming, and window closure when inference raisesAlso updated:
test_qwen_image_pipeline.py(numeric-range coverage through the realforward),test_visual_gen_params.pyandtest_executor_shared_tensor_ipc.py(run_inferenceentry point).Verification:
116 passedacrosstest_profiler,test_qwen_image_pipeline,test_visual_gen_params,test_flux_infer,test_qwen_image_infer,test_executor_shared_tensor_ipc. Run in the nightly staging release container against the merged tree.test_ltx2_pipeline.pywas not run — it requiresLLM_MODELS_ROOT.TLLM_PROFILE_VISUAL_GEN_START_STOP=all, producing a valid 47.7 MB / 128,681-event trace containing text-encoder ops, 15,250 GPU kernels, and VAE-decode convolutions. A separate five-step Nsight run exited 0 and observed the same capture gates.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.