From d787ad9564df1167d01c1be21b676f2d63fc019b Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Thu, 23 Jul 2026 15:09:32 -0700
Subject: [PATCH 1/4] [None][feat] Enable PyTorch profiler traces for VisualGen

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 docs/source/developer-guide/perf-analysis.md  | 13 +++
 tensorrt_llm/_torch/visual_gen/pipeline.py    | 44 +++++++++-
 .../_torch/visual_gen/test_profiler.py        | 88 +++++++++++++++++++
 3 files changed, 144 insertions(+), 1 deletion(-)
 create mode 100644 tests/unittest/_torch/visual_gen/test_profiler.py

diff --git a/docs/source/developer-guide/perf-analysis.md b/docs/source/developer-guide/perf-analysis.md
index e19d32808ea7..d33ea5ff3765 100644
--- a/docs/source/developer-guide/perf-analysis.md
+++ b/docs/source/developer-guide/perf-analysis.md
@@ -55,6 +55,19 @@ Append “python-gil” to Nsys “-t” option.
 1. Set environment variable `TLLM_PROFILE_START_STOP=A-B` to specify the range of the iterations to be collected.
 2. Set environment variable `TLLM_TORCH_PROFILE_TRACE=<path>`, and the results will be saved to `<path>`.
 
+For VisualGen, use `TLLM_PROFILE_VISUAL_GEN_START_STOP` instead. Numeric ranges select
+per-request denoise steps, while `predenoise`, `postdenoise`, and `all` select the
+corresponding generation phases. For example:
+
+```bash
+TLLM_PROFILE_VISUAL_GEN_START_STOP=0-4 \
+TLLM_TORCH_PROFILE_TRACE=/tmp/visual-gen-trace.json \
+python examples/visual_gen/quickstart_example.py
+```
+
+Each process writes its trace to a rank-specific path such as
+`/tmp/visual-gen-trace-rank-0.json`.
+
 ### Visualize the PyTorch profiler results
 
 Use [chrome://tracing/](chrome://tracing/) to inspect the saved profile.
diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py
index 13cbe77bdf12..faef706839dd 100644
--- a/tensorrt_llm/_torch/visual_gen/pipeline.py
+++ b/tensorrt_llm/_torch/visual_gen/pipeline.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import itertools
 import os
 import time
@@ -19,6 +22,9 @@
 from .cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig, SharedGraphPool
 from .modules.vae.parallel_vae_interface import ParallelVAEFactory
 
+PROFILE_START_STOP_ENV_VAR_NAME = "TLLM_PROFILE_VISUAL_GEN_START_STOP"
+PROFILE_TRACE_ENV_VAR_NAME = "TLLM_TORCH_PROFILE_TRACE"
+
 
 class ExtraParamSchema(StrictBaseModel):
     """Schema for a model-specific extra parameter.
@@ -74,7 +80,7 @@ def _parse_profile_range():
        either ``stop`` or ``stop-shutdown`` works. For multi-request capture,
        use a numeric range with ``--capture-range-end=repeat:N``.
     """
-    val = os.environ.get("TLLM_PROFILE_VISUAL_GEN_START_STOP")
+    val = os.environ.get(PROFILE_START_STOP_ENV_VAR_NAME)
     if not val:
         return None
     val = val.strip()
@@ -140,6 +146,9 @@ def __init__(self, pipeline_config: "DiffusionPipelineConfig"):
         # CUDA profiler scoping (TLLM_PROFILE_VISUAL_GEN_START_STOP env var)
         self._profile_range = _parse_profile_range()
         self._profiling_active: bool = False
+        self._torch_profiler = None
+        self._torch_profile_trace_path: Optional[str] = None
+        self._setup_torch_profiler()
         # Single-shot guards for predenoise/postdenoise modes — fire once
         # around the first non-warmup denoise() invocation, then disarm.
         self._predenoise_pending: bool = self._profile_range == "predenoise"
@@ -153,10 +162,38 @@ def __init__(self, pipeline_config: "DiffusionPipelineConfig"):
         # graphed transformer.forward if should_compute == True.
         self._setup_cuda_graphs()
 
+    def _setup_torch_profiler(self) -> None:
+        """Configure PyTorch tracing for the VisualGen profiler range."""
+        torch_trace_path = os.environ.get(PROFILE_TRACE_ENV_VAR_NAME)
+        if not torch_trace_path:
+            return
+        if self._profile_range is None:
+            logger.warning(
+                f"{PROFILE_START_STOP_ENV_VAR_NAME} environment variable "
+                "needs to be set to enable the torch trace. Example to profile "
+                f"denoise steps 0-4: export {PROFILE_START_STOP_ENV_VAR_NAME}=0-4"
+            )
+            return
+
+        trace_base, trace_ext = os.path.splitext(torch_trace_path)
+        self._torch_profile_trace_path = f"{trace_base}-rank-{self.rank}{trace_ext}"
+        activities = [
+            torch.profiler.ProfilerActivity.CPU,
+            torch.profiler.ProfilerActivity.CUDA,
+            torch.profiler.ProfilerActivity.XPU,
+        ]
+        self._torch_profiler = torch.profiler.profile(
+            activities=activities,
+            record_shapes=True,
+            with_modules=True,
+        )
+
     def _cuda_profiler_start(self):
         """Start CUDA profiler if configured and not already active."""
         if self._profile_range is not None and not self._profiling_active:
             torch.cuda.cudart().cudaProfilerStart()
+            if self._torch_profiler is not None:
+                self._torch_profiler.start()
             self._profiling_active = True
             if self.rank == 0:
                 logger.info("CUDA profiler started")
@@ -164,6 +201,11 @@ def _cuda_profiler_start(self):
     def _cuda_profiler_stop(self):
         """Stop CUDA profiler if currently active."""
         if self._profiling_active:
+            if self._torch_profiler is not None:
+                self._torch_profiler.stop()
+                self._torch_profiler.export_chrome_trace(self._torch_profile_trace_path)
+                if self.rank == 0:
+                    logger.info(f"PyTorch profiler trace saved to {self._torch_profile_trace_path}")
             torch.cuda.cudart().cudaProfilerStop()
             self._profiling_active = False
             if self.rank == 0:
diff --git a/tests/unittest/_torch/visual_gen/test_profiler.py b/tests/unittest/_torch/visual_gen/test_profiler.py
new file mode 100644
index 000000000000..77d4151150f0
--- /dev/null
+++ b/tests/unittest/_torch/visual_gen/test_profiler.py
@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.visual_gen.pipeline import (
+    PROFILE_START_STOP_ENV_VAR_NAME,
+    PROFILE_TRACE_ENV_VAR_NAME,
+    BasePipeline,
+)
+
+
+def test_setup_torch_profiler(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv(PROFILE_TRACE_ENV_VAR_NAME, "/tmp/visual-gen-trace.json")
+    pipeline = SimpleNamespace(
+        _profile_range=(frozenset({0}), frozenset({4})),
+        _torch_profiler=None,
+        _torch_profile_trace_path=None,
+        rank=2,
+    )
+    torch_profiler = MagicMock()
+
+    with patch.object(torch.profiler, "profile", return_value=torch_profiler) as profile:
+        BasePipeline._setup_torch_profiler(pipeline)
+
+    assert pipeline._torch_profiler is torch_profiler
+    assert pipeline._torch_profile_trace_path == "/tmp/visual-gen-trace-rank-2.json"
+    profile.assert_called_once_with(
+        activities=[
+            torch.profiler.ProfilerActivity.CPU,
+            torch.profiler.ProfilerActivity.CUDA,
+            torch.profiler.ProfilerActivity.XPU,
+        ],
+        record_shapes=True,
+        with_modules=True,
+    )
+
+
+def test_setup_torch_profiler_requires_profile_range(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv(PROFILE_TRACE_ENV_VAR_NAME, "/tmp/visual-gen-trace.json")
+    pipeline = SimpleNamespace(
+        _profile_range=None,
+        _torch_profiler=None,
+        _torch_profile_trace_path=None,
+        rank=0,
+    )
+
+    with (
+        patch.object(torch.profiler, "profile") as profile,
+        patch("tensorrt_llm._torch.visual_gen.pipeline.logger.warning") as warning,
+    ):
+        BasePipeline._setup_torch_profiler(pipeline)
+
+    profile.assert_not_called()
+    warning.assert_called_once()
+    assert PROFILE_START_STOP_ENV_VAR_NAME in warning.call_args.args[0]
+
+
+def test_cuda_profiler_controls_torch_profiler() -> None:
+    torch_profiler = MagicMock()
+    pipeline = SimpleNamespace(
+        _profile_range=(frozenset({0}), frozenset({4})),
+        _profiling_active=False,
+        _torch_profiler=torch_profiler,
+        _torch_profile_trace_path="/tmp/visual-gen-trace-rank-0.json",
+        rank=0,
+    )
+    cudart = MagicMock()
+
+    with (
+        patch("tensorrt_llm._torch.visual_gen.pipeline.torch.cuda.cudart", return_value=cudart),
+        patch("tensorrt_llm._torch.visual_gen.pipeline.logger.info"),
+    ):
+        BasePipeline._cuda_profiler_start(pipeline)
+        BasePipeline._cuda_profiler_stop(pipeline)
+
+    cudart.cudaProfilerStart.assert_called_once_with()
+    torch_profiler.start.assert_called_once_with()
+    torch_profiler.stop.assert_called_once_with()
+    torch_profiler.export_chrome_trace.assert_called_once_with("/tmp/visual-gen-trace-rank-0.json")
+    cudart.cudaProfilerStop.assert_called_once_with()
+    assert not pipeline._profiling_active

From bca61ab658735dd85a90b7a2c9c84a91f144007c Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Thu, 23 Jul 2026 21:37:15 -0700
Subject: [PATCH 2/4] fix: support VisualGen profiling in Qwen-Image

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 docs/source/developer-guide/perf-analysis.md  |  4 ++
 .../models/qwen_image/pipeline_qwen_image.py  |  5 ++
 tensorrt_llm/_torch/visual_gen/pipeline.py    | 64 +++++++++++++------
 .../visual_gen/test_qwen_image_pipeline.py    | 32 ++++++++++
 4 files changed, 85 insertions(+), 20 deletions(-)

diff --git a/docs/source/developer-guide/perf-analysis.md b/docs/source/developer-guide/perf-analysis.md
index d33ea5ff3765..df99250b32eb 100644
--- a/docs/source/developer-guide/perf-analysis.md
+++ b/docs/source/developer-guide/perf-analysis.md
@@ -68,6 +68,10 @@ python examples/visual_gen/quickstart_example.py
 Each process writes its trace to a rank-specific path such as
 `/tmp/visual-gen-trace-rank-0.json`.
 
+These ranges use the existing VisualGen CUDA/Nsight boundaries. Text encoding
+happens before those boundaries and is not included; `all` starts at denoising
+and includes subsequent VAE work through pipeline cleanup.
+
 ### Visualize the PyTorch profiler results
 
 Use [chrome://tracing/](chrome://tracing/) to inspect the saved profile.
diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
index 54c2310cd56d..3eac0bf19a72 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
@@ -485,10 +485,13 @@ def forward(
 
         # Denoise loop.
         timer.mark_denoise_start()
+        self._start_predenoise_profile()
         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)
+        self._start_denoise_profile()
         for i, t in enumerate(timesteps):
+            self._start_step_profile(i)
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
             noise_pred = self.transformer(
                 hidden_states=latents,
@@ -522,7 +525,9 @@ def forward(
             latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
             if latents.dtype != latents_dtype:
                 latents = latents.to(latents_dtype)
+            self._stop_step_profile(i)
 
+        self._start_postdenoise_profile()
         timer.mark_post_start()
         logger.info("Decoding...")
         image = self._decode_latents(latents, height, width)
diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py
index faef706839dd..e4706f3ebd8f 100644
--- a/tensorrt_llm/_torch/visual_gen/pipeline.py
+++ b/tensorrt_llm/_torch/visual_gen/pipeline.py
@@ -58,7 +58,8 @@ def _parse_profile_range():
                            Single-shot.
     * ``postdenoise``    – profile from the end of the last denoise step to
                            pipeline cleanup, covering VAE decode. Single-shot.
-    * ``all``            – profile the full generation forward (denoise + VAE), skip warmup
+    * ``all``            – profile from denoise through pipeline cleanup,
+                           including VAE decode, but not text encoding; skip warmup
     * (unset)            – no profiler API calls; plain ``nsys profile`` captures everything
 
     Returns ``None`` when unset, one of ``"all"`` / ``"predenoise"`` /
@@ -211,6 +212,43 @@ def _cuda_profiler_stop(self):
             if self.rank == 0:
                 logger.info("CUDA profiler stopped")
 
+    def _start_predenoise_profile(self) -> None:
+        """Start the single-shot pre-denoise profiling window."""
+        if self._predenoise_pending and not self._is_warmup:
+            self._cuda_profiler_start()
+
+    def _start_denoise_profile(self) -> None:
+        """Start all-mode profiling and close the pre-denoise window."""
+        if self._profile_range == "all" and not self._is_warmup:
+            self._cuda_profiler_start()
+        if self._predenoise_pending and not self._is_warmup:
+            self._cuda_profiler_stop()
+            self._predenoise_pending = False
+
+    def _start_step_profile(self, step_index: int) -> None:
+        """Start profiling when the denoise step begins a configured range."""
+        if (
+            isinstance(self._profile_range, tuple)
+            and step_index in self._profile_range[0]
+            and not self._is_warmup
+        ):
+            self._cuda_profiler_start()
+
+    def _stop_step_profile(self, step_index: int) -> None:
+        """Stop profiling when the denoise step ends a configured range."""
+        if (
+            isinstance(self._profile_range, tuple)
+            and step_index in self._profile_range[1]
+            and not self._is_warmup
+        ):
+            self._cuda_profiler_stop()
+
+    def _start_postdenoise_profile(self) -> None:
+        """Start the single-shot post-denoise profiling window."""
+        if self._postdenoise_pending and not self._is_warmup:
+            self._cuda_profiler_start()
+            self._postdenoise_pending = False
+
     def _setup_cuda_graphs(self):
         """Wrap all transformer components with CUDA graph capture/replay.
 
@@ -1097,8 +1135,7 @@ def denoise(
         # TeaCache reset) is captured. The window closes at the first step.
         # Note: hooked here (not at warmup() exit) to avoid leaving the profiler
         # on across the worker's IPC idle, which can interact badly with CUPTI.
-        if self._predenoise_pending and not self._is_warmup:
-            self._cuda_profiler_start()
+        self._start_predenoise_profile()
 
         if timesteps is None:
             timesteps = scheduler.timesteps
@@ -1139,20 +1176,10 @@ def denoise(
 
         # CUDA profiler scoping: "all" starts here (covers denoise + VAE),
         # step ranges start/stop at specific indices. See _parse_profile_range().
-        prof = self._profile_range
-        if prof == "all" and not self._is_warmup:
-            self._cuda_profiler_start()
-        # ``predenoise`` was started in warmup() exit; close the window now,
-        # before the first denoise step kernels run. Single-shot: disarm.
-        if self._predenoise_pending and not self._is_warmup:
-            self._cuda_profiler_stop()
-            self._predenoise_pending = False
-        prof_step_starts = prof[0] if isinstance(prof, tuple) else None
-        prof_step_stops = prof[1] if isinstance(prof, tuple) else None
+        self._start_denoise_profile()
 
         for i, t in enumerate(timesteps):
-            if prof_step_starts is not None and i in prof_step_starts and not self._is_warmup:
-                self._cuda_profiler_start()
+            self._start_step_profile(i)
 
             step_start = time.time()
 
@@ -1230,14 +1257,11 @@ def denoise(
                 )
 
             # Step-level profiler stop
-            if prof_step_stops is not None and i in prof_step_stops and not self._is_warmup:
-                self._cuda_profiler_stop()
+            self._stop_step_profile(i)
 
         # ``postdenoise`` mode: arm the profiler now so the VAE decode (and
         # any post-denoise host work) is captured up to cleanup(). Single-shot.
-        if self._postdenoise_pending and not self._is_warmup:
-            self._cuda_profiler_start()
-            self._postdenoise_pending = False
+        self._start_postdenoise_profile()
 
         if self.rank == 0:
             total_time = time.time() - start_time
diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
index d2d3f08bbd87..d6b489370455 100644
--- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
+++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
@@ -10,6 +10,8 @@
 combination) and the non-CFG path, on CPU.
 """
 
+from unittest.mock import MagicMock
+
 import torch
 
 from tensorrt_llm._torch.visual_gen.models.qwen_image import QwenImagePipeline
@@ -81,6 +83,14 @@ def _pipeline_with_test_doubles():
     pipe.vae_scale_factor = 8
     pipe.transformer = _RecordingTransformer()
     pipe.scheduler = _RecordingScheduler()
+    pipe._profile_range = None
+    pipe._profiling_active = False
+    pipe._torch_profiler = None
+    pipe._torch_profile_trace_path = None
+    pipe._is_warmup = False
+    pipe._predenoise_pending = False
+    pipe._postdenoise_pending = False
+    pipe.rank = 0
     captured = {"encoded_prompts": []}
 
     def _encode_prompt(prompt, device, max_sequence_length):
@@ -211,3 +221,25 @@ def test_forward_runs_true_cfg_pipeline():
     )
     assert torch.allclose(pipe.scheduler.step_calls[0]["noise_pred"], expected_cfg_noise)
     assert pipe.scheduler.step_calls[0]["return_dict"] is False
+
+
+def test_forward_honors_profile_step_range():
+    pipe, _ = _pipeline_with_test_doubles()
+    pipe._profile_range = (frozenset({0}), frozenset({1}))
+    pipe._cuda_profiler_start = MagicMock()
+    pipe._cuda_profiler_stop = MagicMock()
+
+    pipe.forward(
+        prompt=["a cat"],
+        negative_prompt=None,
+        height=32,
+        width=48,
+        num_inference_steps=2,
+        true_cfg_scale=1.0,
+        seed=123,
+        max_sequence_length=16,
+        sigmas=[1.0, 0.5],
+    )
+
+    pipe._cuda_profiler_start.assert_called_once_with()
+    pipe._cuda_profiler_stop.assert_called_once_with()

From f930d1a4f33c1f72cd2c53e751baad8f305cd67b Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Thu, 23 Jul 2026 21:39:09 -0700
Subject: [PATCH 3/4] test: verify Qwen-Image torch trace export

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 .../visual_gen/test_qwen_image_pipeline.py    | 40 +++++++++++--------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
index d6b489370455..28ee0e14087c 100644
--- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
+++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
@@ -10,7 +10,7 @@
 combination) and the non-CFG path, on CPU.
 """
 
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import torch
 
@@ -226,20 +226,28 @@ def test_forward_runs_true_cfg_pipeline():
 def test_forward_honors_profile_step_range():
     pipe, _ = _pipeline_with_test_doubles()
     pipe._profile_range = (frozenset({0}), frozenset({1}))
-    pipe._cuda_profiler_start = MagicMock()
-    pipe._cuda_profiler_stop = MagicMock()
+    pipe._torch_profiler = MagicMock()
+    pipe._torch_profile_trace_path = "visual-gen-trace-rank-0.json"
+    cudart = MagicMock()
 
-    pipe.forward(
-        prompt=["a cat"],
-        negative_prompt=None,
-        height=32,
-        width=48,
-        num_inference_steps=2,
-        true_cfg_scale=1.0,
-        seed=123,
-        max_sequence_length=16,
-        sigmas=[1.0, 0.5],
-    )
+    with patch(
+        "tensorrt_llm._torch.visual_gen.pipeline.torch.cuda.cudart",
+        return_value=cudart,
+    ):
+        pipe.forward(
+            prompt=["a cat"],
+            negative_prompt=None,
+            height=32,
+            width=48,
+            num_inference_steps=2,
+            true_cfg_scale=1.0,
+            seed=123,
+            max_sequence_length=16,
+            sigmas=[1.0, 0.5],
+        )
 
-    pipe._cuda_profiler_start.assert_called_once_with()
-    pipe._cuda_profiler_stop.assert_called_once_with()
+    cudart.cudaProfilerStart.assert_called_once_with()
+    pipe._torch_profiler.start.assert_called_once_with()
+    pipe._torch_profiler.stop.assert_called_once_with()
+    pipe._torch_profiler.export_chrome_trace.assert_called_once_with("visual-gen-trace-rank-0.json")
+    cudart.cudaProfilerStop.assert_called_once_with()

From 5db6e01fc3f3354068abc3ff312fd4595271fd02 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Thu, 23 Jul 2026 23:35:01 -0700
Subject: [PATCH 4/4] fix: broaden VisualGen profiler coverage

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 docs/source/developer-guide/perf-analysis.md  |  13 +-
 tensorrt_llm/_torch/visual_gen/executor.py    |   6 +-
 .../models/qwen_image/pipeline_qwen_image.py  |   5 +-
 tensorrt_llm/_torch/visual_gen/pipeline.py    | 144 ++++++++++-------
 .../integration/test_lists/test-db/l0_a10.yml |   1 +
 .../_torch/visual_gen/test_profiler.py        | 145 +++++++++++++++++-
 .../visual_gen/test_qwen_image_pipeline.py    |  11 +-
 .../visual_gen/test_visual_gen_params.py      |   4 +-
 .../test_executor_shared_tensor_ipc.py        |   3 +
 9 files changed, 254 insertions(+), 78 deletions(-)

diff --git a/docs/source/developer-guide/perf-analysis.md b/docs/source/developer-guide/perf-analysis.md
index df99250b32eb..091934ffb6df 100644
--- a/docs/source/developer-guide/perf-analysis.md
+++ b/docs/source/developer-guide/perf-analysis.md
@@ -66,11 +66,14 @@ python examples/visual_gen/quickstart_example.py
 ```
 
 Each process writes its trace to a rank-specific path such as
-`/tmp/visual-gen-trace-rank-0.json`.
-
-These ranges use the existing VisualGen CUDA/Nsight boundaries. Text encoding
-happens before those boundaries and is not included; `all` starts at denoising
-and includes subsequent VAE work through pipeline cleanup.
+`/tmp/visual-gen-trace-rank-0.json`. If a process captures more than one
+window, later traces add a window suffix such as
+`/tmp/visual-gen-trace-rank-0-window-1.json`.
+
+These ranges use the existing VisualGen CUDA/Nsight boundaries. `all` captures
+the complete request from text encoding through VAE decode. `predenoise`
+captures text encoding, latent preparation, and denoise-loop setup;
+`postdenoise` captures VAE decode and the remaining request work.
 
 ### Visualize the PyTorch profiler results
 
diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py
index 95a170120795..74bda9a0e25b 100644
--- a/tensorrt_llm/_torch/visual_gen/executor.py
+++ b/tensorrt_llm/_torch/visual_gen/executor.py
@@ -235,7 +235,7 @@ class DiffusionRequest:
     (a :class:`~tensorrt_llm.visual_gen.params.VisualGenParams` instance).
     When ``params`` is ``None`` (the default), the executor creates a
     ``VisualGenParams()`` and fills it with pipeline-specific defaults
-    before calling ``pipeline.infer()``.
+    before calling ``pipeline.run_inference()``.
     """
 
     request_id: int
@@ -429,12 +429,12 @@ def process_request(self, req: DiffusionRequest):
                     f"torch.compile recompilation or CUDA graph capture. "
                     f"Warmed-up shapes: {self.pipeline._warmed_up_shapes}"
                 )
-            # Host wall-clock around pipeline.infer(). The pipeline already
+            # Host wall-clock around pipeline.run_inference(). The pipeline already
             # syncs at the end (decode_latents path), so this captures the
             # full executor-side envelope including any pre/post-pipeline work
             # that the per-phase CUDA-event timings on PipelineOutput do not.
             generation_start = time.perf_counter()
-            output = self.pipeline.infer(req)
+            output = self.pipeline.run_inference(req)
             generation = time.perf_counter() - generation_start  # seconds
             if self.rank == 0:
                 # CUDA IPC handles are invalid within the producing process, so
diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
index 3eac0bf19a72..4fda31cd29f5 100644
--- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
+++ b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
@@ -485,11 +485,10 @@ def forward(
 
         # Denoise loop.
         timer.mark_denoise_start()
-        self._start_predenoise_profile()
         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)
-        self._start_denoise_profile()
+        self._profile_denoise_start()
         for i, t in enumerate(timesteps):
             self._start_step_profile(i)
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
@@ -527,7 +526,7 @@ def forward(
                 latents = latents.to(latents_dtype)
             self._stop_step_profile(i)
 
-        self._start_postdenoise_profile()
+        self._profile_denoise_end()
         timer.mark_post_start()
         logger.info("Decoding...")
         image = self._decode_latents(latents, height, width)
diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py
index e4706f3ebd8f..c2de7bcbd6f2 100644
--- a/tensorrt_llm/_torch/visual_gen/pipeline.py
+++ b/tensorrt_llm/_torch/visual_gen/pipeline.py
@@ -52,14 +52,13 @@ def _parse_profile_range():
     * ``A-B``            – profile denoise steps A through B
     * ``A-B,C-D,...``    – multiple ranges; profiler toggles on/off per range
     * ``A,B,...``        – individual steps treated as single-step ranges
-    * ``predenoise``     – profile the per-request pre-loop work inside
-                           ``denoise()`` (CFG config setup, scheduler refresh,
-                           TeaCache reset) up to the first denoise step.
-                           Single-shot.
+    * ``predenoise``     – profile from request start through denoise-loop
+                           setup, including text encoding and latent
+                           preparation. Single-shot.
     * ``postdenoise``    – profile from the end of the last denoise step to
-                           pipeline cleanup, covering VAE decode. Single-shot.
-    * ``all``            – profile from denoise through pipeline cleanup,
-                           including VAE decode, but not text encoding; skip warmup
+                           request completion, covering VAE decode. Single-shot.
+    * ``all``            – profile the full request, from text encoding through
+                           VAE decode; skip warmup
     * (unset)            – no profiler API calls; plain ``nsys profile`` captures everything
 
     Returns ``None`` when unset, one of ``"all"`` / ``"predenoise"`` /
@@ -67,7 +66,7 @@ def _parse_profile_range():
     for numeric ranges.
 
     .. note::
-       Step indices are **per-request**: each ``denoise()`` call resets the
+       Step indices are **per-request**: each denoise loop resets the
        loop counter to 0, so e.g. ``0-4`` profiles steps 0-4 of *every*
        request. This differs from the LLM path's ``TLLM_PROFILE_START_STOP``
        which indexes a global executor iteration counter (one forward pass
@@ -75,11 +74,10 @@ def _parse_profile_range():
 
        ``predenoise`` and ``postdenoise`` are **single-shot per process**:
        they fire once around the first user request after warmup and do not
-       re-arm on subsequent requests. Pair ``predenoise`` with
+       re-arm on subsequent requests. Pair either mode with
        ``nsys --capture-range-end=stop`` (keeps the app running cleanly after
-       collection ends). ``postdenoise`` ends collection at process exit, so
-       either ``stop`` or ``stop-shutdown`` works. For multi-request capture,
-       use a numeric range with ``--capture-range-end=repeat:N``.
+       collection ends). ``all`` and numeric ranges re-arm for each request;
+       use ``--capture-range-end=repeat:N`` to collect multiple requests.
     """
     val = os.environ.get(PROFILE_START_STOP_ENV_VAR_NAME)
     if not val:
@@ -149,9 +147,10 @@ def __init__(self, pipeline_config: "DiffusionPipelineConfig"):
         self._profiling_active: bool = False
         self._torch_profiler = None
         self._torch_profile_trace_path: Optional[str] = None
+        self._torch_profile_window: int = 0
         self._setup_torch_profiler()
         # Single-shot guards for predenoise/postdenoise modes — fire once
-        # around the first non-warmup denoise() invocation, then disarm.
+        # around the first non-warmup request, then disarm.
         self._predenoise_pending: bool = self._profile_range == "predenoise"
         self._postdenoise_pending: bool = self._profile_range == "postdenoise"
 
@@ -178,52 +177,92 @@ def _setup_torch_profiler(self) -> None:
 
         trace_base, trace_ext = os.path.splitext(torch_trace_path)
         self._torch_profile_trace_path = f"{trace_base}-rank-{self.rank}{trace_ext}"
-        activities = [
-            torch.profiler.ProfilerActivity.CPU,
-            torch.profiler.ProfilerActivity.CUDA,
-            torch.profiler.ProfilerActivity.XPU,
-        ]
+        self._create_torch_profiler()
+
+    def _create_torch_profiler(self) -> None:
+        """Create a fresh PyTorch profiler for the next capture window."""
+        activities = [torch.profiler.ProfilerActivity.CPU]
+        if torch.cuda.is_available():
+            activities.append(torch.profiler.ProfilerActivity.CUDA)
+        elif (
+            hasattr(torch, "xpu")
+            and torch.xpu.is_available()
+            and hasattr(torch.profiler.ProfilerActivity, "XPU")
+        ):
+            activities.append(torch.profiler.ProfilerActivity.XPU)
         self._torch_profiler = torch.profiler.profile(
             activities=activities,
             record_shapes=True,
-            with_modules=True,
         )
 
-    def _cuda_profiler_start(self):
+    def _torch_profile_output_path(self) -> str:
+        """Return a non-overwriting path for the current capture window."""
+        if self._torch_profile_trace_path is None:
+            raise RuntimeError("PyTorch profiler trace path is not configured")
+        if self._torch_profile_window == 0:
+            return self._torch_profile_trace_path
+        trace_base, trace_ext = os.path.splitext(self._torch_profile_trace_path)
+        return f"{trace_base}-window-{self._torch_profile_window}{trace_ext}"
+
+    def _cuda_profiler_start(self) -> None:
         """Start CUDA profiler if configured and not already active."""
         if self._profile_range is not None and not self._profiling_active:
-            torch.cuda.cudart().cudaProfilerStart()
-            if self._torch_profiler is not None:
-                self._torch_profiler.start()
+            if self._torch_profile_trace_path is not None and self._torch_profiler is None:
+                self._create_torch_profiler()
+            cudart = torch.cuda.cudart()
+            try:
+                cudart.cudaProfilerStart()
+                if self._torch_profiler is not None:
+                    self._torch_profiler.start()
+            except RuntimeError:
+                cudart.cudaProfilerStop()
+                self._torch_profiler = None
+                raise
             self._profiling_active = True
             if self.rank == 0:
                 logger.info("CUDA profiler started")
 
-    def _cuda_profiler_stop(self):
+    def _cuda_profiler_stop(self) -> None:
         """Stop CUDA profiler if currently active."""
         if self._profiling_active:
-            if self._torch_profiler is not None:
-                self._torch_profiler.stop()
-                self._torch_profiler.export_chrome_trace(self._torch_profile_trace_path)
-                if self.rank == 0:
-                    logger.info(f"PyTorch profiler trace saved to {self._torch_profile_trace_path}")
-            torch.cuda.cudart().cudaProfilerStop()
-            self._profiling_active = False
+            torch_profiler = self._torch_profiler
+            try:
+                if torch_profiler is not None:
+                    torch_profiler.stop()
+                    trace_path = self._torch_profile_output_path()
+                    torch_profiler.export_chrome_trace(trace_path)
+                    self._torch_profile_window += 1
+                    if self.rank == 0:
+                        logger.info(f"PyTorch profiler trace saved to {trace_path}")
+            finally:
+                self._torch_profiler = None
+                try:
+                    torch.cuda.cudart().cudaProfilerStop()
+                finally:
+                    self._profiling_active = False
             if self.rank == 0:
                 logger.info("CUDA profiler stopped")
 
-    def _start_predenoise_profile(self) -> None:
-        """Start the single-shot pre-denoise profiling window."""
-        if self._predenoise_pending and not self._is_warmup:
-            self._cuda_profiler_start()
-
-    def _start_denoise_profile(self) -> None:
-        """Start all-mode profiling and close the pre-denoise window."""
-        if self._profile_range == "all" and not self._is_warmup:
-            self._cuda_profiler_start()
-        if self._predenoise_pending and not self._is_warmup:
+    def run_inference(self, req: Any) -> Any:
+        """Run model-specific inference within shared request profiler boundaries."""
+        profile_request = not self._is_warmup
+        try:
+            if profile_request and self._profile_range == "all":
+                self._cuda_profiler_start()
+            elif profile_request and self._predenoise_pending:
+                self._predenoise_pending = False
+                self._cuda_profiler_start()
+            return self.infer(req)
+        finally:
+            if profile_request:
+                # Ends all/postdenoise windows and safely closes a numeric
+                # range if model inference raises before its normal stop.
+                self._cuda_profiler_stop()
+
+    def _profile_denoise_start(self) -> None:
+        """Close the pre-denoise window at the denoise-loop boundary."""
+        if self._profile_range == "predenoise" and not self._is_warmup:
             self._cuda_profiler_stop()
-            self._predenoise_pending = False
 
     def _start_step_profile(self, step_index: int) -> None:
         """Start profiling when the denoise step begins a configured range."""
@@ -243,7 +282,7 @@ def _stop_step_profile(self, step_index: int) -> None:
         ):
             self._cuda_profiler_stop()
 
-    def _start_postdenoise_profile(self) -> None:
+    def _profile_denoise_end(self) -> None:
         """Start the single-shot post-denoise profiling window."""
         if self._postdenoise_pending and not self._is_warmup:
             self._cuda_profiler_start()
@@ -1130,13 +1169,6 @@ def denoise(
             Single latents if no extra_streams
             Tuple (primary_latents, extra_streams_dict) if extra_streams provided
         """
-        # ``predenoise`` mode: arm the profiler at the very start of denoise()
-        # so the per-request pre-loop work (CFG config, scheduler refresh,
-        # TeaCache reset) is captured. The window closes at the first step.
-        # Note: hooked here (not at warmup() exit) to avoid leaving the profiler
-        # on across the worker's IPC idle, which can interact badly with CUPTI.
-        self._start_predenoise_profile()
-
         if timesteps is None:
             timesteps = scheduler.timesteps
 
@@ -1174,9 +1206,9 @@ def denoise(
 
         start_time = time.time()
 
-        # CUDA profiler scoping: "all" starts here (covers denoise + VAE),
-        # step ranges start/stop at specific indices. See _parse_profile_range().
-        self._start_denoise_profile()
+        # Close the request-level pre-denoise window before the first step.
+        # Numeric ranges start/stop at specific indices.
+        self._profile_denoise_start()
 
         for i, t in enumerate(timesteps):
             self._start_step_profile(i)
@@ -1259,9 +1291,9 @@ def denoise(
             # Step-level profiler stop
             self._stop_step_profile(i)
 
-        # ``postdenoise`` mode: arm the profiler now so the VAE decode (and
-        # any post-denoise host work) is captured up to cleanup(). Single-shot.
-        self._start_postdenoise_profile()
+        # ``postdenoise`` mode: arm the profiler now so the VAE decode and
+        # remaining request work are captured. run_inference() closes it.
+        self._profile_denoise_end()
 
         if self.rank == 0:
             total_time = time.time() - start_time
diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml
index f142378a7687..12de1d81052a 100644
--- a/tests/integration/test_lists/test-db/l0_a10.yml
+++ b/tests/integration/test_lists/test-db/l0_a10.yml
@@ -124,6 +124,7 @@ l0_a10:
   - test_e2e.py::test_trtllm_bench_request_rate_and_concurrency[enable_concurrency-]
   - test_e2e.py::test_trtllm_bench_invalid_token_pytorch[TinyLlama-1.1B-Chat-v1.0-TinyLlama-1.1B-Chat-v1.0]
   # visual_gen
+  - unittest/_torch/visual_gen/test_profiler.py
   - unittest/_torch/visual_gen/test_visual_gen_params.py
   - unittest/visual_gen/test_output.py
   - unittest/visual_gen/test_iteration_stats.py
diff --git a/tests/unittest/_torch/visual_gen/test_profiler.py b/tests/unittest/_torch/visual_gen/test_profiler.py
index 77d4151150f0..eeeec0427272 100644
--- a/tests/unittest/_torch/visual_gen/test_profiler.py
+++ b/tests/unittest/_torch/visual_gen/test_profiler.py
@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -22,9 +22,13 @@ def test_setup_torch_profiler(monkeypatch: pytest.MonkeyPatch) -> None:
         _torch_profile_trace_path=None,
         rank=2,
     )
+    pipeline._create_torch_profiler = MethodType(BasePipeline._create_torch_profiler, pipeline)
     torch_profiler = MagicMock()
 
-    with patch.object(torch.profiler, "profile", return_value=torch_profiler) as profile:
+    with (
+        patch.object(torch.cuda, "is_available", return_value=True),
+        patch.object(torch.profiler, "profile", return_value=torch_profiler) as profile,
+    ):
         BasePipeline._setup_torch_profiler(pipeline)
 
     assert pipeline._torch_profiler is torch_profiler
@@ -33,10 +37,8 @@ def test_setup_torch_profiler(monkeypatch: pytest.MonkeyPatch) -> None:
         activities=[
             torch.profiler.ProfilerActivity.CPU,
             torch.profiler.ProfilerActivity.CUDA,
-            torch.profiler.ProfilerActivity.XPU,
         ],
         record_shapes=True,
-        with_modules=True,
     )
 
 
@@ -69,8 +71,12 @@ def test_cuda_profiler_controls_torch_profiler() -> None:
         _profiling_active=False,
         _torch_profiler=torch_profiler,
         _torch_profile_trace_path="/tmp/visual-gen-trace-rank-0.json",
+        _torch_profile_window=0,
         rank=0,
     )
+    pipeline._torch_profile_output_path = MethodType(
+        BasePipeline._torch_profile_output_path, pipeline
+    )
     cudart = MagicMock()
 
     with (
@@ -86,3 +92,134 @@ def test_cuda_profiler_controls_torch_profiler() -> None:
     torch_profiler.export_chrome_trace.assert_called_once_with("/tmp/visual-gen-trace-rank-0.json")
     cudart.cudaProfilerStop.assert_called_once_with()
     assert not pipeline._profiling_active
+    assert pipeline._torch_profiler is None
+
+
+def test_cuda_profiler_uses_fresh_trace_for_each_window() -> None:
+    first_profiler = MagicMock()
+    second_profiler = MagicMock()
+    pipeline = SimpleNamespace(
+        _profile_range=(frozenset({0, 2}), frozenset({1, 3})),
+        _profiling_active=False,
+        _torch_profiler=first_profiler,
+        _torch_profile_trace_path="/tmp/visual-gen-trace-rank-0.json",
+        _torch_profile_window=0,
+        rank=0,
+    )
+    pipeline._create_torch_profiler = MethodType(BasePipeline._create_torch_profiler, pipeline)
+    pipeline._torch_profile_output_path = MethodType(
+        BasePipeline._torch_profile_output_path, pipeline
+    )
+    cudart = MagicMock()
+
+    with (
+        patch.object(torch.profiler, "profile", return_value=second_profiler) as profile,
+        patch("tensorrt_llm._torch.visual_gen.pipeline.torch.cuda.cudart", return_value=cudart),
+        patch("tensorrt_llm._torch.visual_gen.pipeline.logger.info"),
+    ):
+        BasePipeline._cuda_profiler_start(pipeline)
+        BasePipeline._cuda_profiler_stop(pipeline)
+        BasePipeline._cuda_profiler_start(pipeline)
+        BasePipeline._cuda_profiler_stop(pipeline)
+
+    profile.assert_called_once()
+    first_profiler.export_chrome_trace.assert_called_once_with("/tmp/visual-gen-trace-rank-0.json")
+    second_profiler.export_chrome_trace.assert_called_once_with(
+        "/tmp/visual-gen-trace-rank-0-window-1.json"
+    )
+    assert cudart.cudaProfilerStart.call_count == 2
+    assert cudart.cudaProfilerStop.call_count == 2
+    assert pipeline._torch_profile_window == 2
+
+
+def test_cuda_profiler_closes_cuda_gate_when_trace_export_fails() -> None:
+    torch_profiler = MagicMock()
+    torch_profiler.export_chrome_trace.side_effect = RuntimeError("export failed")
+    pipeline = SimpleNamespace(
+        _profiling_active=True,
+        _torch_profiler=torch_profiler,
+        _torch_profile_trace_path="/tmp/visual-gen-trace-rank-0.json",
+        _torch_profile_window=0,
+        rank=0,
+    )
+    pipeline._torch_profile_output_path = MethodType(
+        BasePipeline._torch_profile_output_path, pipeline
+    )
+    cudart = MagicMock()
+
+    with (
+        patch("tensorrt_llm._torch.visual_gen.pipeline.torch.cuda.cudart", return_value=cudart),
+        pytest.raises(RuntimeError, match="export failed"),
+    ):
+        BasePipeline._cuda_profiler_stop(pipeline)
+
+    cudart.cudaProfilerStop.assert_called_once_with()
+    assert not pipeline._profiling_active
+    assert pipeline._torch_profiler is None
+
+
+class _RequestProfilingPipeline:
+    def __init__(self, profile_range: str) -> None:
+        self._profile_range = profile_range
+        self._profiling_active = False
+        self._is_warmup = False
+        self._predenoise_pending = profile_range == "predenoise"
+        self._postdenoise_pending = profile_range == "postdenoise"
+        self.events: list[str] = []
+
+    def _cuda_profiler_start(self) -> None:
+        if not self._profiling_active:
+            self.events.append("start")
+            self._profiling_active = True
+
+    def _cuda_profiler_stop(self) -> None:
+        if self._profiling_active:
+            self.events.append("stop")
+            self._profiling_active = False
+
+    def infer(self, req: object) -> object:
+        self.events.append("text_encode")
+        BasePipeline._profile_denoise_start(self)
+        self.events.append("denoise")
+        BasePipeline._profile_denoise_end(self)
+        self.events.append("vae_decode")
+        return req
+
+
+@pytest.mark.parametrize(
+    ("profile_range", "expected_events"),
+    [
+        ("all", ["start", "text_encode", "denoise", "vae_decode", "stop"]),
+        ("predenoise", ["start", "text_encode", "stop", "denoise", "vae_decode"]),
+        ("postdenoise", ["text_encode", "denoise", "start", "vae_decode", "stop"]),
+    ],
+)
+def test_run_inference_owns_request_profile_boundaries(
+    profile_range: str, expected_events: list[str]
+) -> None:
+    pipeline = _RequestProfilingPipeline(profile_range)
+    request = object()
+
+    output = BasePipeline.run_inference(pipeline, request)
+
+    assert output is request
+    assert pipeline.events == expected_events
+
+
+def test_run_inference_defers_predenoise_profile_until_after_warmup() -> None:
+    pipeline = _RequestProfilingPipeline("predenoise")
+    pipeline._is_warmup = True
+
+    BasePipeline.run_inference(pipeline, object())
+
+    assert pipeline.events == ["text_encode", "denoise", "vae_decode"]
+    assert pipeline._predenoise_pending
+    assert not pipeline._profiling_active
+
+    pipeline._is_warmup = False
+    pipeline.events.clear()
+    BasePipeline.run_inference(pipeline, object())
+
+    assert pipeline.events == ["start", "text_encode", "stop", "denoise", "vae_decode"]
+    assert not pipeline._predenoise_pending
+    assert not pipeline._profiling_active
diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
index 28ee0e14087c..1f4f690c991b 100644
--- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
+++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py
@@ -87,10 +87,10 @@ def _pipeline_with_test_doubles():
     pipe._profiling_active = False
     pipe._torch_profiler = None
     pipe._torch_profile_trace_path = None
+    pipe._torch_profile_window = 0
     pipe._is_warmup = False
     pipe._predenoise_pending = False
     pipe._postdenoise_pending = False
-    pipe.rank = 0
     captured = {"encoded_prompts": []}
 
     def _encode_prompt(prompt, device, max_sequence_length):
@@ -226,7 +226,8 @@ def test_forward_runs_true_cfg_pipeline():
 def test_forward_honors_profile_step_range():
     pipe, _ = _pipeline_with_test_doubles()
     pipe._profile_range = (frozenset({0}), frozenset({1}))
-    pipe._torch_profiler = MagicMock()
+    torch_profiler = MagicMock()
+    pipe._torch_profiler = torch_profiler
     pipe._torch_profile_trace_path = "visual-gen-trace-rank-0.json"
     cudart = MagicMock()
 
@@ -247,7 +248,7 @@ def test_forward_honors_profile_step_range():
         )
 
     cudart.cudaProfilerStart.assert_called_once_with()
-    pipe._torch_profiler.start.assert_called_once_with()
-    pipe._torch_profiler.stop.assert_called_once_with()
-    pipe._torch_profiler.export_chrome_trace.assert_called_once_with("visual-gen-trace-rank-0.json")
+    torch_profiler.start.assert_called_once_with()
+    torch_profiler.stop.assert_called_once_with()
+    torch_profiler.export_chrome_trace.assert_called_once_with("visual-gen-trace-rank-0.json")
     cudart.cudaProfilerStop.assert_called_once_with()
diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py
index e69c81d21f1b..d47eb9ec1f5c 100644
--- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py
+++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py
@@ -1132,7 +1132,7 @@ def test_caller_params_not_mutated(self):
 class TestEngineFailureTransport:
     """Validation is enforced at :meth:`VisualGen.generate_async` entry, so
     by the time a request reaches ``process_request`` only runtime
-    failures from ``pipeline.infer()`` can produce an error response.
+    failures from ``pipeline.run_inference()`` can produce an error response.
     The error message rides back on ``DiffusionResponse.error_msg``.
     """
 
@@ -1162,7 +1162,7 @@ def test_runtime_error_carried_on_response(self):
         executor._merge_defaults = lambda req: DiffusionExecutor._merge_defaults(executor, req)
         executor.pipeline.warmup_cache_key = MagicMock(return_value=(1024, 1024, None))
         executor.pipeline._warmed_up_shapes = None
-        executor.pipeline.infer = MagicMock(side_effect=RuntimeError("oops"))
+        executor.pipeline.run_inference = MagicMock(side_effect=RuntimeError("oops"))
 
         req = DiffusionRequest(
             request_id=7,
diff --git a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py
index 29fce319cdc7..d53762badfb2 100644
--- a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py
+++ b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py
@@ -122,6 +122,9 @@ def warmup_cache_key(self, height, width, num_frames):
     def infer(self, req):
         return PipelineOutput(video=_expected_video(), frame_rate=24.0)
 
+    def run_inference(self, req):
+        return self.infer(req)
+
     def cleanup(self):
         pass