[None][feat] Enable PyTorch profiler traces for VisualGen

docs/source/developer-guide/perf-analysis.md

@@ -55,6 +55,26 @@ 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`. 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
 
 Use [chrome://tracing/](chrome://tracing/) to inspect the saved profile.

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

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

@@ -488,7 +488,9 @@ def forward(
         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._profile_denoise_start()
         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 +524,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._profile_denoise_end()
         timer.mark_post_start()
         logger.info("Decoding...")
         image = self._decode_latents(latents, height, width)

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.
@@ -46,35 +52,34 @@ 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 the full generation forward (denoise + VAE), 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"`` /
     ``"postdenoise"`` for keyword modes, or ``(frozenset(starts), frozenset(stops))``
     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
        services all in-flight requests, so there is no "per request" index).
 
        ``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("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,8 +145,12 @@ 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._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"
 
@@ -153,22 +162,132 @@ def __init__(self, pipeline_config: "DiffusionPipelineConfig"):
         # graphed transformer.forward if should_compute == True.
         self._setup_cuda_graphs()
 
-    def _cuda_profiler_start(self):
+    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}"
+        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,
+        )
+
+    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_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:
-            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 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()
+
+    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 _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()
+            self._postdenoise_pending = False
+
     def _setup_cuda_graphs(self):
         """Wrap all transformer components with CUDA graph capture/replay.
 
@@ -1050,14 +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.
-        if self._predenoise_pending and not self._is_warmup:
-            self._cuda_profiler_start()
-
         if timesteps is None:
             timesteps = scheduler.timesteps
 
@@ -1095,22 +1206,12 @@ 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().
-        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
+        # 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):
-            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()
 
@@ -1188,14 +1289,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
+        # ``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

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