[TRTLLM-15447][feat] Add model startup stage timing metrics

docs/source/llm-api/index.md

@@ -50,44 +50,143 @@ llm = LLM(model=<local_path_to_model>)
 
 ## Startup Metrics
 
-For the PyTorch backend, the beta `LLM.startup_metrics` property reports weight-loading timings from
-worker rank 0. Values are wall-clock seconds. The property returns an empty dictionary when the
-backend does not provide startup metrics.
+For the PyTorch backend, the beta `LLM.startup_metrics` property reports executor construction,
+model-engine warmup, and weight-loading timings from worker rank 0. Values are wall-clock seconds.
+The property returns an empty dictionary when the backend does not provide startup metrics.
 
 ```python
 llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 print(llm.startup_metrics)
 ```
 
-A typical result has the following structure:
+A typical result has the following structure (values are illustrative):
 
 ```json
 {
+  "initial_model_engine": {
+    "sampling_warmup_seconds": 0.342,
+    "attention_warmup_seconds": 2.911,
+    "general_warmup_seconds": 0.387,
+    "autotuner_warmup_seconds": 0.041,
+    "mamba_hybrid_warmup_seconds": 0.000,
+    "ctx_cuda_graph_warmup_seconds": 0.480,
+    "ctx_cuda_graph_capture_seconds": 0.320,
+    "post_ctx_cuda_graph_capture_warmup_seconds": 0.110,
+    "gen_cuda_graph_warmup_seconds": 0.566,
+    "gen_cuda_graph_capture_seconds": 0.950,
+    "dg_paged_mqa_warmup_seconds": 0.000,
+    "cute_dsl_radix_topk_warmup_seconds": 0.000,
+    "memory_pool_prepopulation_seconds": 0.094,
+    "kv_cache_cleanup_seconds": 0.194,
+    "total_warmup_seconds": 5.485
+  },
+  "final_model_engine": {
+    "sampling_warmup_seconds": 0.291,
+    "attention_warmup_seconds": 0.057,
+    "general_warmup_seconds": 0.492,
+    "autotuner_warmup_seconds": 0.046,
+    "mamba_hybrid_warmup_seconds": 0.000,
+    "gen_cuda_graph_warmup_seconds": 0.599,
+    "gen_cuda_graph_capture_seconds": 0.991,
+    "dg_paged_mqa_warmup_seconds": 0.000,
+    "cute_dsl_radix_topk_warmup_seconds": 0.000,
+    "memory_pool_prepopulation_seconds": 0.097,
+    "kv_cache_cleanup_seconds": 1.657,
+    "total_warmup_seconds": 4.230
+  },
+  "py_executor": {
+    "config_and_checkpoint_loader_initialization_seconds": 0.004,
+    "model_engine_creation_seconds": 1.554,
+    "sampler_creation_seconds": 1.195,
+    "initial_kv_cache_creation_seconds": 0.025,
+    "speculative_decoding_resource_manager_creation_seconds": 0.000,
+    "speculative_drafter_creation_seconds": 0.000,
+    "initial_py_executor_creation_seconds_for_kv_cache_estimation": 5.766,
+    "kv_cache_capacity_configuration_seconds": 0.075,
+    "final_kv_cache_creation_seconds": 0.057,
+    "final_py_executor_creation_seconds": 4.459,
+    "worker_start_seconds": 0.000,
+    "total_py_executor_creation_seconds": 13.373
+  },
   "model_loader": {
-    "total_model_loading_seconds":  1.971,
-    "checkpoint_preparation_seconds": 1.177,
-    "weight_population_seconds": 0.598,
+    "checkpoint_preparation_seconds": 0.719,
+    "weight_population_seconds": 0.272,
     "checkpoint_finalization_seconds": 0.031,
-    "post_load_processing_seconds": 0.005
+    "post_load_processing_seconds": 0.004,
+    "total_model_loading_seconds": 1.545
   }
 }
 ```
 
-The `model_loader` object contains timings for the main LLM weights. If a draft
-model is used, the additional fields `draft_checkpoint_preparation_seconds`,
-`draft_weight_population_seconds`, and `draft_checkpoint_finalization_seconds`
-appear.
+Some of the timing metrics may not appear or may show as a very small value if the corresponding codepath is not executed, e.g.
+if the prefill CUDA graphs are not enabled, its related metrics "...ctx_cuda_graph..." will not appear.
+
+The `py_executor` property contains the timed scopes in `create_py_executor()`, which creates the [PyExecutor](../developer-guide/overview.md).
+
+| PyExecutor metric | Scope |
+|-----------------|-------|
+| `config_and_checkpoint_loader_initialization_seconds` | Initialize the model config and checkpoint loader. |
+| `model_engine_creation_seconds` | Construct the main model engine, which includes loading model weights. |
+| `draft_model_engine_creation_seconds` | Construct the separate draft model engine in the deprecated two-model MTP setting. |
+| `guided_decoder_creation_seconds` | Construct guided-decoding resources. |
+| `sampler_creation_seconds` | Construct the sampler. |
+| `initial_kv_cache_creation_seconds` | Construct the temporary KV cache used to estimate final capacity. |
+| `speculative_decoding_resource_manager_creation_seconds` | Construct the speculative-decoding resource manager. |
+| `speculative_drafter_creation_seconds` | Construct the speculative drafter. |
+| `initial_py_executor_creation_seconds_for_kv_cache_estimation` | Construct the temporary PyExecutor for KV cache capacity estimation, which includes the initial model engine warmup. |
+| `kv_cache_capacity_configuration_seconds` | Determine final KV cache capacity from the temporary executor and available memory. |
+| `final_kv_cache_creation_seconds` | Construct the final KV cache retained for serving. |
+| `final_py_executor_creation_seconds` | Using the final KV cache, construct the final PyExecutor retained for serving, which includes the final model engine warmup. |
+| `worker_start_seconds` | Start the final PyExecutor worker. |
+| `total_py_executor_creation_seconds` | Total time for the `create_py_executor()` call, including all applicable scopes above. |
+
+The two model-engine properties contain times for various `PyTorchModelEngine` warmup stages.
+`initial_model_engine` measures the timings for the initial model engine creation as part of the
+`py_executor.initial_py_executor_creation_seconds_for_kv_cache_estimation` timing.
+Likewise, `final_model_engine` measures the timings for the final model engine creation as part of
+the `py_executor.final_py_executor_creation_seconds` timing.
+The deprecated two-model speculative decoding setting may produce the additional properties
+`initial_draft_model_engine` and `final_draft_model_engine`.
+
+Each model-engine property contains the warmup scopes below. Which scopes appear depends on the model
+and configuration; for example, encoder-decoder and context-parallel configurations can skip some
+stages.
+
+| Model-engine metric | Scope |
+|---------------------|-------|
+| `sampling_warmup_seconds` | Warm up sampling kernels and run `torch.compile()` on sampling operators. |
+| `attention_warmup_seconds` | Warm up the attention backend and kernels. |
+| `general_warmup_seconds` | Prepare and warm up general input shapes and release temporary workspaces. |
+| `autotuner_warmup_seconds` | Run kernel autotuning warmup. |
+| `mamba_hybrid_warmup_seconds` | Warm up Mamba hybrid kernels, when applicable. |
+| `gen_cuda_graph_warmup_seconds` | Run the warmup-only pass for generation and mixed encoder-decoder CUDA graphs, including LoRA autotuning and pipeline-parallel cache synchronization and cleanup. |
+| `gen_cuda_graph_capture_seconds` | Capture generation and mixed encoder-decoder CUDA graphs for serving. |
+| `ctx_cuda_graph_warmup_seconds` | Warm up prefill forwards before CUDA graph capture. |
+| `ctx_cuda_graph_capture_seconds` | Capture prefill CUDA graphs. |
+| `post_ctx_cuda_graph_capture_warmup_seconds` | Warm up prefill logits buffers with many requests after CUDA graph capture. |
+| `dg_paged_mqa_warmup_seconds` | Warm up DeepGEMM paged-MQA metadata, when applicable. |
+| `cute_dsl_radix_topk_warmup_seconds` | Warm up the CuTe DSL radix top-k kernel, when applicable. |
+| `memory_pool_prepopulation_seconds` | Prepare maximum-shape requests and pre-populate the memory pool. |
+| `kv_cache_cleanup_seconds` | Check and clear invalid KV cache values produced during warmup. |
+| `total_warmup_seconds` | Complete model-engine warmup, including KV cache cleanup. |
+
+The `model_loader` property contains timings for loading the main LLM weights, including
+`checkpoint_finalization_seconds` for checkpoint-session synchronization and cleanup after weight population.
+If a draft model is used, additional fields `draft_checkpoint_preparation_seconds`,
+`draft_weight_population_seconds`, and `draft_checkpoint_finalization_seconds` will appear.
+A `draft_model_loader` property can also appear in the deprecated two-model MTP setting.
+The `py_executor.model_engine_creation_seconds` timing includes the total `model_loader` timing.
 
 | Metric | Description |
 |--------|-------------|
-| `total_model_loading_seconds` | Overall model construction and loading interval measured after checkpoint configuration validation. It includes the named phases below. |
 | `checkpoint_preparation_seconds` | Time spent warming up, parsing and preparing checkpoint tensors for the model. Some checkpoint formats can populate model storage directly during this phase. |
 | `weight_population_seconds` | Time spent copying prepared checkpoint tensors into model parameters on GPUs. This metric can be absent for formats that populate weights directly during the above checkpoint preparation phase. |
-| `checkpoint_finalization_seconds` | Time spent finalizing the checkpoint session after weight population. This includes loader-specific synchronization and cleanup; rank-striped read-ahead includes waiting for peer ranks and stopping background readers. |
+| `checkpoint_finalization_seconds` | Time spent finalizing the checkpoint session after weight population. This includes loader-specific synchronization and cleanup. |
 | `draft_checkpoint_preparation_seconds` | Checkpoint preparation time for draft weights loaded as part of the model loader. |
 | `draft_weight_population_seconds` | Weight population time for draft weights loaded as part of the model loader. |
 | `draft_checkpoint_finalization_seconds` | Checkpoint finalization time for draft weights loaded as part of the model loader. |
 | `post_load_processing_seconds` | Time spent in format-specific hooks and model finalization, including post-load weight transformation, quantization and memory cleanup. |
+| `total_model_loading_seconds` | Overall model construction and loading interval measured after checkpoint configuration validation. It includes the named phases below. |
 
 `trtllm-serve` exposes the same rank-0 payload in the `startup_metrics` field of the
 `GET /server_info` response:
@@ -100,7 +199,7 @@ curl http://localhost:8000/server_info
 {
   "startup_metrics": {
     "model_loader": {
-      "total_model_loading_seconds": 1.971,
+      "total_model_loading_seconds": 1.545,
       ...
     }
   }

tensorrt_llm/_startup.py

@@ -49,11 +49,19 @@ def mark_initialization(self, name: str) -> None:
         logger.info(f"[startup][pid={os.getpid()}] {self.name}/{name}: done in {elapsed:.3f}s")
 
     @contextmanager
-    def phase(self, name: str) -> Iterator[None]:
+    def phase(
+        self,
+        name: str,
+        *,
+        metrics: dict[str, float] | None = None,
+        metric_name: str | None = None,
+    ) -> Iterator[None]:
         """Record outermost elapsed time, including failures, and log status.
 
         Failed work still contributes to startup time; its exception propagates
         and its completion marker is explicitly labelled ``failed``.
+        When ``metrics`` is provided, accumulate the same elapsed value under
+        ``metric_name`` (or ``name``), including nested phases.
         """
         logger.info(f"[startup][pid={os.getpid()}] {self.name}/{name}: start")
         start = time.perf_counter()
@@ -68,6 +76,9 @@ def phase(self, name: str) -> Iterator[None]:
             self.depth -= 1
             if outermost:
                 self.timings[name] = self.timings.get(name, 0.0) + elapsed
+            if metrics is not None:
+                key = metric_name if metric_name is not None else name
+                metrics[key] = metrics.get(key, 0.0) + elapsed
             status = "done" if completed else "failed"
             logger.info(
                 f"[startup][pid={os.getpid()}] {self.name}/{name}: {status} in {elapsed:.3f}s"

tensorrt_llm/_torch/compilation/piecewise_optimizer.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import dataclasses
 from typing import Callable, List, Optional, Sequence, Union
 from unittest.mock import patch
@@ -144,6 +147,8 @@ class Entry:
 
 class PiecewiseRunner(object):
 
+    WARMUP_STEPS = 3
+
     def __init__(
         self,
         graph: GraphModule,
@@ -223,7 +228,7 @@ def __call__(self, *args):
             if not get_capture_piecewise_cuda_graph_flag():
                 return entry.callable(*args)
 
-            if entry.warmup_count < 3:
+            if entry.warmup_count < self.WARMUP_STEPS:
                 entry.warmup_count += 1
                 return entry.callable(*args)
 

tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py

@@ -8,6 +8,7 @@
 import torch
 from torch import nn
 
+from ..models.modeling_utils import timing_metric
 from ..utils import make_weak_ref
 from .breakable_cuda_graph import (
     BreakableCUDAGraph,
@@ -28,6 +29,9 @@ class BreakableCUDAGraphRunner:
 
     _WARMUP_STEPS = 2
 
+    CUDA_GRAPH_WARMUP_METRIC = "cuda_graph_warmup_seconds"
+    CUDA_GRAPH_CAPTURE_METRIC = "cuda_graph_capture_seconds"
+
     def __init__(self, layer_model: nn.Module) -> None:
         self.layer_model = layer_model
         self._graphs: dict[int, BreakableCUDAGraph] = {}
@@ -38,6 +42,12 @@ def __init__(self, layer_model: nn.Module) -> None:
         self._state = BreakableCUDAGraphRunnerState.IDLE
         self._active_graph: Optional[BreakableCUDAGraph] = None
         self._active_num_tokens: Optional[int] = None
+        self._metrics: dict[str, float] = {}
+
+    @property
+    def metrics(self) -> dict[str, float]:
+        """Warmup and capture timings for the most recent token bucket."""
+        return self._metrics
 
     @property
     def state(self) -> BreakableCUDAGraphRunnerState:
@@ -67,7 +77,12 @@ def warmup(self, engine_forward: Callable[[], Any], steps: int = _WARMUP_STEPS)
             self._state = BreakableCUDAGraphRunnerState.IDLE
 
     def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None:
-        """Warm up eagerly, then capture one prefill token bucket."""
+        """Warm up eagerly, then capture one prefill token bucket.
+
+        Record each phase in metrics, replacing the previous bucket's timings.
+        Both phases run on the capture stream and synchronize before timing ends.
+        """
+        self._metrics.clear()
         if self._state != BreakableCUDAGraphRunnerState.IDLE:
             raise RuntimeError(f"Cannot capture BCG while runner is {self._state.value}")
         if num_tokens in self._graphs:
@@ -79,21 +94,25 @@ def capture(self, num_tokens: int, engine_forward: Callable[[], Any]) -> None:
         created_memory_pool = False
         try:
             with torch.cuda.stream(self._capture_stream):
-                self.warmup(engine_forward)
-
-                # Every segment in the first BCG bucket must receive the same
-                # explicit pool handle. Passing None lets each CUDAGraph create
-                # its own private pool, which multiplies the model workspace by
-                # the number of eager breaks.
-                if self._memory_pool is None:
-                    self._memory_pool = torch.cuda.graph_pool_handle()
-                    created_memory_pool = True
-
-                self._state = BreakableCUDAGraphRunnerState.CAPTURE
-                graph = BreakableCUDAGraph()
-                self._active_graph = graph
-                self._active_num_tokens = num_tokens
-                output = engine_forward()
+                with timing_metric(self.CUDA_GRAPH_WARMUP_METRIC, self._metrics):
+                    self.warmup(engine_forward)
+                    torch.cuda.synchronize()
+
+                with timing_metric(self.CUDA_GRAPH_CAPTURE_METRIC, self._metrics):
+                    # Every segment in the first BCG bucket must receive the same
+                    # explicit pool handle. Passing None lets each CUDAGraph create
+                    # its own private pool, which multiplies the model workspace by
+                    # the number of eager breaks.
+                    if self._memory_pool is None:
+                        self._memory_pool = torch.cuda.graph_pool_handle()
+                        created_memory_pool = True
+
+                    self._state = BreakableCUDAGraphRunnerState.CAPTURE
+                    graph = BreakableCUDAGraph()
+                    self._active_graph = graph
+                    self._active_num_tokens = num_tokens
+                    output = engine_forward()
+                    torch.cuda.synchronize()
 
             current_stream.wait_stream(self._capture_stream)
             if not torch.is_tensor(output):