NVIDIA · chang-l · Aug 1, 2026 · Jul 23, 2026 · Jul 24, 2026 · Jul 24, 2026
@@ -55,6 +55,37 @@ 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.
+
+Two contracts worth knowing: a numeric range never extends past the denoise loop
+it selects (a stop index beyond the last step closes at the last step instead),
+and on a pipeline that runs more than one denoise loop per request, the per-loop
+modes apply to each loop — a numeric range writes one trace per loop, while
+`postdenoise` arms after the first loop rather than the last. All windows open at
+the pipeline's inference entry point, so executor-side request preparation falls
+outside them even though it counts toward the reported `generation` latency.
+
+For the per-mode specifics, see `parse_profile_range` in
+`tensorrt_llm/_torch/visual_gen/profiler.py`.
+
 ### Visualize the PyTorch profiler results
 
 Use [chrome://tracing/](chrome://tracing/) to inspect the saved profile.

@@ -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
@@ -444,7 +444,7 @@ def process_request(self, req: DiffusionRequest):
                     f"torch.compile recompilation or CUDA graph capture. "
                     f"Warmed-up shapes: {self.pipeline._warmed_up_shapes}"
                 )
-            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

@@ -1722,7 +1722,11 @@ def _refinement_denoise(
             timer.mark_stage2_start()
 
         # --- Euler denoising loop (no guidance) ---
-        for i in range(len(sigmas) - 1):
+        # Stage 2 runs its own loop rather than BasePipeline.denoise(), so it
+        # drives the profiler windows itself. Without this a numeric range
+        # would capture stage 1 only and the trace would silently omit half
+        # the denoising.
+        for i, _ in self._profile_denoise_steps(range(len(sigmas) - 1)):
             with nvtx_range(f"refinement_step {i}"):
                 sigma = sigmas[i]
                 sigma_next = sigmas[i + 1]

@@ -568,7 +568,7 @@ def forward(
         timer.mark_denoise_start()
         logger.info("Denoising (%d steps)...", len(timesteps))
         cuda_graph_enabled = self.pipeline_config.cuda_graph.enable
-        for i, t in enumerate(timesteps):
+        for i, t in self._profile_denoise_steps(timesteps):
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
             if do_cfg_parallel:
                 local_embeds, local_mask = self._select_cfg_inputs(

@@ -488,7 +488,7 @@ def forward(
 
         timer.mark_denoise_start()
         logger.info("Denoising edit (%d steps)...", len(timesteps))
-        for t in timesteps:
+        for _, t in self._profile_denoise_steps(timesteps):
             latent_model_input = torch.cat([latents, image_latents], dim=1)
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
 

@@ -907,7 +907,7 @@ def forward(
         additional_t_cond = torch.zeros(batch_size, device=device, dtype=torch.long)
         timer.mark_denoise_start()
         logger.info("Denoising layered output (%d steps)...", len(timesteps))
-        for t in timesteps:
+        for _, t in self._profile_denoise_steps(timesteps):
             latent_model_input = torch.cat([latents, image_latents], dim=1)
             timestep = t.expand(latents.shape[0]).to(latents.dtype)
             noise_pred = self.transformer(

@@ -1,8 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import contextlib
 import itertools
 import os
 import time
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+)
 
 import torch
 import torch.distributed as dist
@@ -21,6 +36,7 @@
 from .cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig, SharedGraphPool
 from .mapping import _VisualGenAutotuneDist
 from .modules.vae.parallel_vae_interface import ParallelVAEFactory
+from .profiler import VisualGenProfiler
 
 
 class ExtraParamSchema(StrictBaseModel):
@@ -38,67 +54,6 @@ class ExtraParamSchema(StrictBaseModel):
     )
 
 
-def _parse_profile_range():
-    """Parse ``TLLM_PROFILE_VISUAL_GEN_START_STOP`` for CUDA profiler scoping.
-
-    Visual-gen-specific env var (separate from the LLM path's
-    ``TLLM_PROFILE_START_STOP``). Use with ``nsys profile -c cudaProfilerApi ...``.
-
-    Supported formats:
-
-    * ``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.
-    * ``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
-    * (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
-       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
-       ``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``.
-    """
-    val = os.environ.get("TLLM_PROFILE_VISUAL_GEN_START_STOP")
-    if not val:
-        return None
-    val = val.strip()
-    if val.lower() in ("all", "predenoise", "postdenoise"):
-        return val.lower()
-    # Parse comma-separated ranges: "A-B,C-D,..." or single steps "A,B,..."
-    # Same format as the LLM path (PyExecutor._load_iteration_indexes).
-    starts, stops = [], []
-    for span in val.split(","):
-        span = span.strip()
-        if "-" in span:
-            start, stop = span.split("-", 1)
-            starts.append(int(start))
-            stops.append(int(stop))
-        else:
-            v = int(span)
-            starts.append(v)
-            stops.append(v)
-    return frozenset(starts), frozenset(stops)
-
-
 if TYPE_CHECKING:
     from .cache import CacheAccelerator
     from .config import DiffusionPipelineConfig
@@ -140,13 +95,8 @@ def __init__(self, pipeline_config: "DiffusionPipelineConfig"):
         self.scheduler: Optional[Any] = None
         self._is_warmup: bool = False
 
-        # CUDA profiler scoping (TLLM_PROFILE_VISUAL_GEN_START_STOP env var)
-        self._profile_range = _parse_profile_range()
-        self._profiling_active: bool = False
-        # 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"
-        self._postdenoise_pending: bool = self._profile_range == "postdenoise"
+        # Profiler window scoping (TLLM_PROFILE_VISUAL_GEN_START_STOP env var)
+        self._profiler = VisualGenProfiler(rank=self.rank)
 
         # Initialize transformer
         self._init_transformer()
@@ -156,21 +106,27 @@ def __init__(self, pipeline_config: "DiffusionPipelineConfig"):
         # graphed transformer.forward if should_compute == True.
         self._setup_cuda_graphs()
 
-    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()
-            self._profiling_active = True
-            if self.rank == 0:
-                logger.info("CUDA profiler started")
+    def run_inference(self, req: Any) -> Any:
+        """Run model-specific inference within shared request profiler boundaries."""
+        if self._is_warmup:
+            return self.infer(req)
+        with self._profiler.request_scope():
+            return self.infer(req)
 
-    def _cuda_profiler_stop(self):
-        """Stop CUDA profiler if currently active."""
-        if self._profiling_active:
-            torch.cuda.cudart().cudaProfilerStop()
-            self._profiling_active = False
-            if self.rank == 0:
-                logger.info("CUDA profiler stopped")
+    def _profile_denoise_steps(self, timesteps: Iterable[Any]) -> Iterator[Tuple[int, Any]]:
+        """Enumerate a denoise loop's steps within its profiling windows.
+
+        ``BasePipeline.denoise()`` uses this, and so must any pipeline that
+        writes its own denoise loop (Qwen-Image, LTX-2 stage 2) — otherwise
+        that loop is silently absent from every trace. Substitute it for the
+        loop's ``enumerate()``; it owns every window boundary the loop has::
+
+            for i, t in self._profile_denoise_steps(timesteps):
+                ...
+        """
+        if self._is_warmup:
+            return enumerate(timesteps)
+        return self._profiler.steps(timesteps)
 
     def _setup_cuda_graphs(self):
         """Wrap all transformer components with CUDA graph capture/replay.
@@ -1129,14 +1085,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
 
@@ -1174,23 +1122,7 @@ 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
-
-        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()
-
+        for i, t in self._profile_denoise_steps(timesteps):
             step_start = time.time()
 
             current_guidance_scale = self._resolve_step_guidance_scale(
@@ -1267,16 +1199,6 @@ def denoise(
                     f"Avg={avg_time:.2f}s/step ETA={eta:.1f}s"
                 )
 
-            # 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()
-
-        # ``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
-
         if self.rank == 0:
             total_time = time.time() - start_time
             logger.info("=" * 80)
@@ -1309,7 +1231,7 @@ def denoise(
 
     def cleanup(self):
         """Call before dist.destroy_process_group()."""
-        self._cuda_profiler_stop()
+        self._profiler.close_window()
 
         for name, runner in self._cuda_graph_runners.items():
             logger.info(f"Releasing CUDA graphs for {name}")