From f07b455a5d364896286c11c823c83ab6c00f0284 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Tue, 4 Aug 2026 23:10:15 -0700
Subject: [PATCH 01/23] [TRTLLM-15447][feat] Add model startup stage timing
 metrics

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md                  | 54 +++++++++--
 .../_torch/pyexecutor/model_engine.py         | 66 ++++++++-----
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  6 ++
 .../_torch/pyexecutor/py_executor_creator.py  | 92 +++++++++++++++++--
 tensorrt_llm/executor/base_worker.py          | 13 +++
 ...utor_creator_flash_mla_tokens_per_block.py |  2 +-
 ...y_executor_creator_mla_cache_reuse_sync.py | 77 +++++++++++++++-
 .../test_pytorch_model_engine_warmup.py       | 20 ++++
 tests/unittest/executor/test_base_worker.py   | 59 ++++++++++++
 9 files changed, 351 insertions(+), 38 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 2d0a9f0cf64f..6ce9fef9f5d4 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -50,9 +50,9 @@ 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")
@@ -63,6 +63,18 @@ A typical result has the following structure:
 
 ```json
 {
+  "executor": {
+    "total_executor_creation_seconds": 12.341,
+    "model_engine_creation_seconds": 5.241,
+    "final_kv_cache_creation_seconds": 0.421,
+    "final_py_executor_creation_seconds": 4.317
+  },
+  "final_model_engine": {
+    "total_warmup_seconds": 4.217,
+    "attention_warmup_seconds": 0.301,
+    "autotuner_warmup_seconds": 1.732,
+    "cuda_graph_capture_seconds": 1.105
+  },
   "model_loader": {
     "total_model_loading_seconds":  1.971,
     "checkpoint_preparation_seconds": 1.177,
@@ -73,10 +85,38 @@ A typical result has the following structure:
 }
 ```
 
-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.
+The `executor` object contains the timed scopes in `create_py_executor()`. When KV cache capacity
+estimation creates a temporary executor before constructing the final executor,
+`initial_model_engine` contains the warmup metrics from the temporary pass and
+`final_model_engine` contains the warmup metrics from the final pass. Without estimation, only
+`final_model_engine` appears. The model-engine dictionaries are promoted alongside `executor` in
+the startup metrics payload so callers do not need to traverse another nesting level. Legacy
+two-model speculative decoding can likewise produce `initial_draft_model_engine` and
+`final_draft_model_engine`.
+
+Executor scope names describe the constructed component rather than the internal memory-category
+enum. The initial KV cache and PyExecutor scopes are temporary resources used for KV cache capacity
+estimation; their final counterparts are the resources retained for serving.
+
+| Executor metric | Scope |
+|-----------------|-------|
+| `model_engine_creation_seconds` | Construct and load the main model engine. |
+| `draft_model_engine_creation_seconds` | Construct and load the separate draft model engine, when used. |
+| `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. |
+| `final_kv_cache_creation_seconds` | Construct the final KV cache retained for serving. |
+| `speculative_decoding_resource_manager_creation_seconds` | Construct speculative-decoding resource managers. |
+| `speculative_drafter_creation_seconds` | Construct the speculative drafter. |
+| `initial_py_executor_creation_seconds_for_kv_cache_estimation` | Construct and warm up the temporary PyExecutor used during KV cache capacity estimation. |
+| `kv_cache_capacity_configuration_seconds` | Determine final KV cache capacity from the temporary executor and available memory. |
+| `final_py_executor_creation_seconds` | Construct and warm up the final PyExecutor retained for serving. |
+| `worker_start_seconds` | Start the final executor worker. |
+| `total_executor_creation_seconds` | Complete `create_py_executor()` call, including all applicable scopes above. |
+
+The `model_loader` object contains timings for the main LLM weights. If a draft model is used,
+additional field `draft_checkpoint_preparation_seconds` and `draft_weight_population_seconds` will appear.
+A `draft_model_loader` object can also appear in deprecated 2-model style MTP setting.
 
 | Metric | Description |
 |--------|-------------|
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index da57e8c27eda..7a155559075b 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -55,7 +55,7 @@
 from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader
 from ..models.modeling_multimodal_mixin import (MultimodalModelMixin,
                                                 _build_request_multimodal_input)
-from ..models.modeling_utils import DecoderModelForCausalLM
+from ..models.modeling_utils import DecoderModelForCausalLM, timing_metric
 from ..modules.mamba.mamba2_metadata import Mamba2Metadata
 from ..moe.expert_statistic import ExpertStatistic
 from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer,
@@ -378,6 +378,7 @@ def __init__(
     ):
         _configure_deep_gemm_pdl()
 
+        self._metrics: dict[str, float] = {}
         self.forward_pass_callable = None
         self._cleanup_done = False
         self._model_caller: Optional[ModelCaller] = None
@@ -1180,6 +1181,11 @@ def moe_load_balancer_iter_info(self, value: Tuple[bool, bool]):
     def use_beam_search(self):
         return self.max_beam_width > 1
 
+    @property
+    def metrics(self) -> dict[str, float]:
+        """Return model-engine warmup time metrics."""
+        return self._metrics
+
     def _get_draft_kv_cache_manager(
         self, resource_manager: ResourceManager
     ) -> Optional[Union[KVCacheManager, KVCacheManagerV2]]:
@@ -1296,16 +1302,17 @@ def wrapper(self,
                     *args,
                     **kwargs):
             result = method(self, resource_manager, *args, **kwargs)
-            kv_cache_manager = (resource_manager.get_resource_manager(
-                self.kv_cache_manager_key)
-                                if resource_manager is not None else None)
-            if kv_cache_manager is not None:
-                has_invalid_values = kv_cache_manager.check_invalid_values_in_kv_cache(
-                    fill_with_zero=True)
-                if has_invalid_values:
-                    logger.warning(
-                        "NaNs/Infs have been introduced to KVCache during warmup, KVCache was filled with zeros to avoid potential issues"
-                    )
+            with timing_metric("kv_cache_cleanup_seconds", self._metrics):
+                kv_cache_manager = (resource_manager.get_resource_manager(
+                    self.kv_cache_manager_key)
+                                    if resource_manager is not None else None)
+                if kv_cache_manager is not None:
+                    has_invalid_values = kv_cache_manager.check_invalid_values_in_kv_cache(
+                        fill_with_zero=True)
+                    if has_invalid_values:
+                        logger.warning(
+                            "NaNs/Infs have been introduced to KVCache during warmup, KVCache was filled with zeros to avoid potential issues"
+                        )
             return result
 
         return wrapper
@@ -1420,8 +1427,12 @@ def forward_encoder(
 
     @with_warmup_flag
     @warmup_with_kv_cache_cleanup
-    def warmup(self,
-               resource_manager: Optional[ResourceManager] = None) -> None:
+    def warmup(self, resource_manager: Optional[ResourceManager] = None) -> None:
+        """Run model warmup and record its total wall-clock duration."""
+        with timing_metric("total_warmup_seconds", self._metrics):
+            self._warmup_impl(resource_manager)
+
+    def _warmup_impl(self, resource_manager: Optional[ResourceManager] = None) -> None:
         """
         Orchestrates the warmup process by calling specialized warmup methods for
         torch.compile, the autotuner, and CUDA graphs.
@@ -1497,7 +1508,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             self._prewarm_cute_dsl_indexer_q()
         log_mem_snapshot("warmup/after_cute_dsl_indexer_q")
         if not is_enc_dec:
-            with self._warmup_timer.phase("attention_jit"):
+            with self._warmup_timer.phase("attention_jit"), timing_metric(
+                    "attention_warmup_seconds", self._metrics):
                 self._run_attention_warmup(resource_manager,
                                            can_run_general_warmup)
 
@@ -1507,7 +1519,7 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
                 warmup_requests_configs = self._agree_warmup_shapes(
                     self._get_full_general_warmup_requests(resource_manager))
                 # Currently graph has not been captured, disable cuda graph for this warmup.
-                with self.no_cuda_graph():
+                with timing_metric("general_warmup_seconds", self._metrics), self.no_cuda_graph():
                     self._general_warmup(resource_manager,
                                          warmup_requests_configs)
                     # Release C++ MoE workspace buffers so the autotuner can
@@ -1522,7 +1534,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
         # Helix CP is decode-only and runs into issues with the
         # autotuner warmup's context requests.
         if not is_enc_dec and not self.mapping.has_cp_helix():
-            with self._warmup_timer.phase("autotuner"):
+            with self._warmup_timer.phase("autotuner"), timing_metric(
+                    "autotuner_warmup_seconds", self._metrics):
                 self._run_autotuner_warmup(resource_manager)
             log_mem_snapshot("warmup/after_autotuner")
             # Pre-JIT Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels
@@ -1530,7 +1543,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             # since MambaHybridCacheManager skips _general_warmup and the
             # default autotuner shape is single-seq / no-initstates. Safe
             # no-op for non-Mamba models.
-            with self._warmup_timer.phase("mamba_hybrid"):
+            with self._warmup_timer.phase("mamba_hybrid"), timing_metric(
+                    "mamba_hybrid_warmup_seconds", self._metrics):
                 self._run_mamba_hybrid_warmup(resource_manager)
             log_mem_snapshot("warmup/after_mamba_hybrid")
             # Release the autotuner's exploration-mode intermediates. The
@@ -1551,12 +1565,15 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             with self.cuda_graph_runner.allow_capture():
                 self.cuda_graph_runner.is_warmup_only = True
                 try:
-                    with self.maybe_autotune_lora():
-                        self._run_cuda_graph_warmup(resource_manager)
+                    with timing_metric("cuda_graph_warmup_seconds",
+                                       self._metrics):
+                        with self.maybe_autotune_lora():
+                            self._run_cuda_graph_warmup(resource_manager)
                 finally:
                     self.cuda_graph_runner.is_warmup_only = False
                 self.cuda_graph_runner.padding_dummy_requests = {}
-                self._run_cuda_graph_warmup(resource_manager)
+                with timing_metric("cuda_graph_capture_seconds", self._metrics):
+                    self._run_cuda_graph_warmup(resource_manager)
         log_mem_snapshot("warmup/after_cuda_graph_capture")
         # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned
         # batch bucket the runtime can produce (max_batch_size scaled by the
@@ -1571,9 +1588,11 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
         # creates; build it when every forward above was skipped.
         with self._warmup_timer.phase("dsa_prewarm"):
             self._ensure_dsa_attn_metadata_for_warmup(resource_manager)
-            self._warmup_dg_paged_mqa_logits_metadata()
+            with timing_metric("dg_paged_mqa_warmup_seconds", self._metrics):
+                self._warmup_dg_paged_mqa_logits_metadata()
             log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata")
-            self._warmup_cute_dsl_radix_topk()
+            with timing_metric("cute_dsl_radix_topk_warmup_seconds", self._metrics):
+                self._warmup_cute_dsl_radix_topk()
         log_mem_snapshot("warmup/after_cute_dsl_radix_topk")
         if can_run_general_warmup:
             # Pre-populate the memory pool with max-shape allocations to reduce
@@ -1581,7 +1600,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             with self._warmup_timer.phase("memory_pool_prepop"):
                 warmup_requests_configs = self._get_max_shape_warmup_requests(
                     resource_manager)
-                self._general_warmup(resource_manager, warmup_requests_configs)
+                with timing_metric("memory_pool_prepopulation_seconds", self._metrics):
+                    self._general_warmup(resource_manager, warmup_requests_configs)
             log_mem_snapshot("warmup/after_memory_pool_prepop")
 
         # Allocate the CUDA graph padding dummies now, while the KV cache is
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 3fc930fa8a48..ef602789237f 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -428,6 +428,7 @@ def __init__(
             dwdp_manager: Optional[DwdpManager] = None,
             enable_kv_pool_rebalance: bool = False):
         super(PyExecutor, self).__init__()
+        self._metrics: dict[str, float | dict[str, float]] = {}
         self.device_id = torch.cuda.current_device()
         self.global_rank = dist.rank
         # Store the execution stream for decoder/model forward operations.
@@ -1078,6 +1079,11 @@ def on_detected():
         if start_worker:
             self.start_worker()
 
+    @property
+    def metrics(self) -> dict[str, float | dict[str, float]]:
+        """Return executor construction and model-engine startup metrics."""
+        return self._metrics
+
     def _maybe_init_kv_connector_manager(self):
         if self.kv_connector_manager is not None:
             if self.kv_cache_transceiver is not None:
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 6204a1cd88bf..edf5b5c48b29 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -35,6 +35,7 @@
 
 from ..attention.backends.interface import AttentionRuntimeFeatures
 from ..distributed import Distributed
+from ..models.modeling_utils import timing_metric
 from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter,
                            get_spec_resource_manager,
                            should_use_separate_draft_kv_cache)
@@ -107,6 +108,31 @@ def _bytes_to_gib(bytes: int) -> float:
         "Draft model for speculative decoding",
     }
 
+    creation_stage_metric_names = {
+        ExecutorMemoryType.SAMPLER:
+        "sampler_creation_seconds",
+        ExecutorMemoryType.DRAFTER:
+        "speculative_drafter_creation_seconds",
+        ExecutorMemoryType.GUIDED_DECODER:
+        "guided_decoder_creation_seconds",
+        ExecutorMemoryType.SPEC_RESOURCES:
+        "speculative_decoding_resource_manager_creation_seconds",
+        ExecutorMemoryType.INIT_KV_CACHE:
+        "initial_kv_cache_creation_seconds",
+        ExecutorMemoryType.INIT_EXTRA_RESOURCES:
+        "initial_py_executor_creation_seconds_for_kv_cache_estimation",
+        ExecutorMemoryType.MODEL_EXTRA:
+        "kv_cache_capacity_configuration_seconds",
+        ExecutorMemoryType.EXTRA_RESOURCES:
+        "final_py_executor_creation_seconds",
+        ExecutorMemoryType.KV_CACHE:
+        "final_kv_cache_creation_seconds",
+        ExecutorMemoryType.MODEL_ENGINE_MAIN:
+        "model_engine_creation_seconds",
+        ExecutorMemoryType.MODEL_ENGINE_DRAFT:
+        "draft_model_engine_creation_seconds",
+    }
+
     # Suggestion to reduce component memory usage
     memory_type_tuning_suggestion = {
         ExecutorMemoryType.SAMPLER:
@@ -346,6 +372,21 @@ def log_memory_usage(stage: str):
     )
 
 
+def _move_model_engine_metrics(
+        py_executor: PyExecutor,
+        model_engine: PyTorchModelEngine,
+        stage: str,
+        draft_model_engine: Optional[PyTorchModelEngine] = None) -> None:
+    """Move model-engine metrics into the executor under a startup stage."""
+    py_executor.metrics[f"{stage}_model_engine"] = dict(model_engine.metrics)
+    model_engine.metrics.clear()
+
+    if draft_model_engine is not None:
+        py_executor.metrics[f"{stage}_draft_model_engine"] = dict(
+            draft_model_engine.metrics)
+        draft_model_engine.metrics.clear()
+
+
 def _create_py_executor_impl(
     llm_args: TorchLlmArgs,
     _startup_timer: _StartupTimer,
@@ -374,10 +415,38 @@ def _create_py_executor_impl(
     Returns:
         A fully initialized PyExecutor instance.
     """
+    creation_metrics: dict[str, float] = {}
+    with timing_metric("total_executor_creation_seconds", creation_metrics):
+        py_executor = _create_py_executor(
+            llm_args=llm_args,
+            _startup_timer=_startup_timer,
+            creation_metrics=creation_metrics,
+            checkpoint_dir=checkpoint_dir,
+            tokenizer=tokenizer,
+            profiling_stage_data=profiling_stage_data,
+            resource_governor_queue=resource_governor_queue,
+        )
+    py_executor.metrics.update(creation_metrics)
+    return py_executor
+
+
+def _create_py_executor(
+    llm_args: TorchLlmArgs,
+    _startup_timer: _StartupTimer,
+    creation_metrics: dict[str, float],
+    checkpoint_dir: Optional[str] = None,
+    tokenizer: Optional[TokenizerBase] = None,
+    profiling_stage_data: Optional[dict] = None,
+    resource_governor_queue=None,
+) -> PyExecutor:
+    """Create and initialize a PyExecutor while recording scoped metrics."""
+    initial_model_engine_metrics: dict[str, float | dict[str, float]] = {}
 
     skip_est = os.environ.get("TRTLLM_SKIP_KV_CACHE_ESTIMATION", '0') == '1'
-    llm_args, checkpoint_loader = _load_config_and_create_checkpoint_loader(
-        llm_args, checkpoint_dir)
+    with timing_metric("config_and_checkpoint_loader_initialization_seconds",
+                       creation_metrics):
+        llm_args, checkpoint_loader = _load_config_and_create_checkpoint_loader(
+            llm_args, checkpoint_dir)
 
     garbage_collection_gen0_threshold = llm_args.garbage_collection_gen0_threshold
     lora_config = llm_args.lora_config
@@ -603,9 +672,12 @@ def allocation_scope(current_stage: ExecutorMemoryType):
             ExecutorMemoryType.KV_CACHE: "final_kv_cache_allocation",
             ExecutorMemoryType.EXTRA_RESOURCES: "final_executor_creation",
         }.get(current_stage, current_stage.value)
-        with _startup_timer.phase(
-                timing_name), mem_monitor.observe_creation_stage(current_stage):
-            stage = current_stage.value
+        stage = current_stage.value
+        metric_name = mem_monitor.creation_stage_metric_names[current_stage]
+        with _startup_timer.phase(timing_name), timing_metric(
+                metric_name,
+                creation_metrics), mem_monitor.observe_creation_stage(
+                    current_stage):
             if not enable_sleep or stage.startswith("_no_capture"):
                 yield
             else:
@@ -1047,6 +1119,9 @@ def allocation_scope(current_stage: ExecutorMemoryType):
 
     if estimating_kv_cache:
         assert kv_cache_creator is not None
+        _move_model_engine_metrics(py_executor, model_engine, "initial",
+                                   draft_model_engine)
+        initial_model_engine_metrics = dict(py_executor.metrics)
         with allocation_scope(ExecutorMemoryType.MODEL_EXTRA):
             kv_cache_creator.configure_kv_cache_capacity(py_executor)
 
@@ -1124,9 +1199,14 @@ def allocation_scope(current_stage: ExecutorMemoryType):
     if mapping.rank == 0:
         logger.info(f"LLM Args:\n{llm_args}")
 
-    with _startup_timer.phase("executor_start_worker"):
+    with _startup_timer.phase("executor_start_worker"), timing_metric(
+            "worker_start_seconds", creation_metrics):
         py_executor.start_worker()
 
+    py_executor.metrics.update(initial_model_engine_metrics)
+    _move_model_engine_metrics(py_executor, model_engine, "final",
+                               draft_model_engine)
+
     return py_executor
 
 
diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py
index c4dfa7552244..c79e7e1c8bcf 100644
--- a/tensorrt_llm/executor/base_worker.py
+++ b/tensorrt_llm/executor/base_worker.py
@@ -1047,6 +1047,19 @@ def get_startup_metrics(self) -> dict:
             return {}
 
         startup_metrics = {}
+        executor_metrics = dict(getattr(self.engine, "metrics", {}))
+        for model_engine_stage in (
+                "initial_model_engine",
+                "final_model_engine",
+                "initial_draft_model_engine",
+                "final_draft_model_engine",
+        ):
+            model_engine_metrics = executor_metrics.pop(model_engine_stage,
+                                                        None)
+            if model_engine_metrics is not None:
+                startup_metrics[model_engine_stage] = dict(model_engine_metrics)
+        startup_metrics["executor"] = executor_metrics
+
         model_engine = getattr(self.engine, "model_engine", None)
         model_loader = getattr(model_engine, "model_loader", None)
         if model_loader is not None:
diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py b/tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py
index 4084bd903639..e248283db6c8 100644
--- a/tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py
+++ b/tests/unittest/_torch/executor/test_py_executor_creator_flash_mla_tokens_per_block.py
@@ -42,7 +42,7 @@
 
 
 def _get_create_py_executor_source() -> str:
-    return inspect.getsource(py_executor_creator.create_py_executor)
+    return inspect.getsource(py_executor_creator._create_py_executor)
 
 
 def test_flash_mla_branch_overrides_local_tokens_per_block():
diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
index ec3276fd6a54..93793c91ea6b 100644
--- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
+++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
@@ -23,13 +24,37 @@
     _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS,
 )
 from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType
-from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, ContextChunkingPolicy
+from tensorrt_llm.llmapi.llm_args import (
+    CacheTransceiverConfig,
+    ContextChunkingPolicy,
+    ExecutorMemoryType,
+)
 from tensorrt_llm.models.modeling_utils import QuantConfig
 from tensorrt_llm.quantization import QuantAlgo
 
 pytestmark = pytest.mark.cpu_only
 
 
+def test_executor_creation_stage_metric_names_are_descriptive():
+    assert py_executor_creator._ExecutorMemoryMonitor.creation_stage_metric_names == {
+        ExecutorMemoryType.SAMPLER: "sampler_creation_seconds",
+        ExecutorMemoryType.DRAFTER: "speculative_drafter_creation_seconds",
+        ExecutorMemoryType.GUIDED_DECODER: "guided_decoder_creation_seconds",
+        ExecutorMemoryType.SPEC_RESOURCES: (
+            "speculative_decoding_resource_manager_creation_seconds"
+        ),
+        ExecutorMemoryType.INIT_KV_CACHE: "initial_kv_cache_creation_seconds",
+        ExecutorMemoryType.INIT_EXTRA_RESOURCES: (
+            "initial_py_executor_creation_seconds_for_kv_cache_estimation"
+        ),
+        ExecutorMemoryType.MODEL_EXTRA: "kv_cache_capacity_configuration_seconds",
+        ExecutorMemoryType.EXTRA_RESOURCES: "final_py_executor_creation_seconds",
+        ExecutorMemoryType.KV_CACHE: "final_kv_cache_creation_seconds",
+        ExecutorMemoryType.MODEL_ENGINE_MAIN: "model_engine_creation_seconds",
+        ExecutorMemoryType.MODEL_ENGINE_DRAFT: "draft_model_engine_creation_seconds",
+    }
+
+
 class _DummyCalibrator:
     """Mock calibrator for testing that bypasses actual calibration logic."""
 
@@ -79,6 +104,7 @@ def __init__(self, resources, model_engine, peft_cache_config, execution_stream)
         """
         self.resource_manager = _DummyResourceManager(resources)
         self.model_engine = model_engine
+        self.metrics = {}
         self.peft_cache_config = peft_cache_config
         self.execution_stream = execution_stream
         self.started = False
@@ -146,6 +172,7 @@ def __init__(
 
         self._warmup_timer = _WarmupTimer(rank=0)
         self.attn_runtime_features = attn_runtime_features
+        self.metrics = {}
         self.max_seq_len = max_seq_len
         self.max_num_tokens = 128
         self.sparse_attention_config = (
@@ -356,6 +383,54 @@ def _create_py_executor_instance(**kwargs):
     )
 
 
+def test_move_model_engine_metrics_moves_and_clears():
+    py_executor = SimpleNamespace(metrics={})
+    model_engine = SimpleNamespace(metrics={"total_warmup_seconds": 2.5})
+    draft_model_engine = SimpleNamespace(metrics={"total_warmup_seconds": 1.5})
+
+    py_executor_creator._move_model_engine_metrics(
+        py_executor, model_engine, "initial", draft_model_engine
+    )
+    model_engine.metrics["total_warmup_seconds"] = 3.5
+    draft_model_engine.metrics["total_warmup_seconds"] = 2.0
+    py_executor_creator._move_model_engine_metrics(
+        py_executor, model_engine, "final", draft_model_engine
+    )
+
+    assert py_executor.metrics == {
+        "initial_model_engine": {"total_warmup_seconds": 2.5},
+        "initial_draft_model_engine": {"total_warmup_seconds": 1.5},
+        "final_model_engine": {"total_warmup_seconds": 3.5},
+        "final_draft_model_engine": {"total_warmup_seconds": 2.0},
+    }
+    assert model_engine.metrics == {}
+    assert draft_model_engine.metrics == {}
+
+
+def test_total_executor_creation_metric_finishes_on_error(monkeypatch):
+    captured_metrics = None
+
+    @contextmanager
+    def _timing_metric(metric_name, metrics):
+        nonlocal captured_metrics
+        captured_metrics = metrics
+        try:
+            yield
+        finally:
+            metrics[metric_name] = 1.5
+
+    def _raise_during_creation(**kwargs):
+        raise RuntimeError("creation failed")
+
+    monkeypatch.setattr(py_executor_creator, "timing_metric", _timing_metric)
+    monkeypatch.setattr(py_executor_creator, "_create_py_executor", _raise_during_creation)
+
+    with pytest.raises(RuntimeError, match="creation failed"):
+        py_executor_creator.create_py_executor(SimpleNamespace())
+
+    assert captured_metrics == {"total_executor_creation_seconds": 1.5}
+
+
 def test_mla_unsupported_sm_fallback_syncs_cache_reuse(monkeypatch):
     """Verify MLA unsupported SM fallback disables cache reuse in both config and runtime.
 
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index edba96b46000..0cb15b701ac8 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -1176,6 +1176,26 @@ def test_native_mxfp8_respects_disabled_global_autotuner(self):
         self.assertFalse(method.use_native_autotuner)
         self.assertFalse(method.needs_native_autotune)
 
+    def test_warmup_stage_metrics_are_recorded(self):
+        model_engine, resource_manager = _build_engine_and_resource_manager()
+        _run_warmup_tracked(model_engine, resource_manager)
+
+        expected_metrics = {
+            "attention_warmup_seconds",
+            "general_warmup_seconds",
+            "autotuner_warmup_seconds",
+            "mamba_hybrid_warmup_seconds",
+            "cuda_graph_warmup_seconds",
+            "cuda_graph_capture_seconds",
+            "dg_paged_mqa_warmup_seconds",
+            "cute_dsl_radix_topk_warmup_seconds",
+            "memory_pool_prepopulation_seconds",
+            "kv_cache_cleanup_seconds",
+            "total_warmup_seconds",
+        }
+        self.assertEqual(expected_metrics, model_engine.metrics.keys())
+        self.assertTrue(all(value >= 0 for value in model_engine.metrics.values()))
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py
index bc9a814e76c7..f8e920402f68 100644
--- a/tests/unittest/executor/test_base_worker.py
+++ b/tests/unittest/executor/test_base_worker.py
@@ -1,4 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
 import time
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -21,6 +24,62 @@
 model_path = llm_models_root() / default_model_name
 
 
+def test_get_startup_metrics_promotes_model_engine_stages():
+    worker = object.__new__(BaseWorker)
+    worker._is_pytorch_backend = True
+    worker.engine = SimpleNamespace(
+        metrics={
+            "worker_start_seconds": 0.5,
+            "initial_model_engine": {
+                "total_warmup_seconds": 2.5
+            },
+            "final_model_engine": {
+                "total_warmup_seconds": 3.5
+            },
+            "initial_draft_model_engine": {
+                "total_warmup_seconds": 1.0
+            },
+            "final_draft_model_engine": {
+                "total_warmup_seconds": 1.25
+            },
+        },
+        model_engine=SimpleNamespace(
+            metrics={"must_not_be_reported": 1.0},
+            model_loader=SimpleNamespace(
+                metrics={"total_model_loading_seconds": 1.5}),
+        ),
+        draft_model_engine=SimpleNamespace(
+            metrics={"must_not_be_reported": 1.0},
+            model_loader=SimpleNamespace(
+                metrics={"total_model_loading_seconds": 0.75}),
+        ),
+    )
+
+    assert worker.get_startup_metrics() == {
+        "initial_model_engine": {
+            "total_warmup_seconds": 2.5
+        },
+        "final_model_engine": {
+            "total_warmup_seconds": 3.5
+        },
+        "initial_draft_model_engine": {
+            "total_warmup_seconds": 1.0
+        },
+        "final_draft_model_engine": {
+            "total_warmup_seconds": 1.25
+        },
+        "executor": {
+            "worker_start_seconds": 0.5
+        },
+        "model_loader": {
+            "total_model_loading_seconds": 1.5
+        },
+        "draft_model_loader": {
+            "total_model_loading_seconds": 0.75
+        },
+    }
+
+
 @pytest.mark.cpu_only
 def test_enqueue_request_wraps_lora_load_error():
 

From 000bd1a6c3231651a9abe3cc6e7494042b6d2228 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Mon, 24 Aug 2026 11:39:28 -0700
Subject: [PATCH 02/23] [TRTLLM-15447][fix] Rename PyExecutor startup metrics
 group

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md                | 12 ++++++------
 tensorrt_llm/executor/base_worker.py        |  2 +-
 tests/unittest/executor/test_base_worker.py |  2 +-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 6ce9fef9f5d4..93751da936e0 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -63,7 +63,7 @@ A typical result has the following structure:
 
 ```json
 {
-  "executor": {
+  "py_executor": {
     "total_executor_creation_seconds": 12.341,
     "model_engine_creation_seconds": 5.241,
     "final_kv_cache_creation_seconds": 0.421,
@@ -85,11 +85,11 @@ A typical result has the following structure:
 }
 ```
 
-The `executor` object contains the timed scopes in `create_py_executor()`. When KV cache capacity
-estimation creates a temporary executor before constructing the final executor,
+The `py_executor` object contains the timed scopes in `create_py_executor()`. When KV cache capacity
+estimation creates a temporary PyExecutor before constructing the final PyExecutor,
 `initial_model_engine` contains the warmup metrics from the temporary pass and
 `final_model_engine` contains the warmup metrics from the final pass. Without estimation, only
-`final_model_engine` appears. The model-engine dictionaries are promoted alongside `executor` in
+`final_model_engine` appears. The model-engine dictionaries are promoted alongside `py_executor` in
 the startup metrics payload so callers do not need to traverse another nesting level. Legacy
 two-model speculative decoding can likewise produce `initial_draft_model_engine` and
 `final_draft_model_engine`.
@@ -98,7 +98,7 @@ Executor scope names describe the constructed component rather than the internal
 enum. The initial KV cache and PyExecutor scopes are temporary resources used for KV cache capacity
 estimation; their final counterparts are the resources retained for serving.
 
-| Executor metric | Scope |
+| PyExecutor metric | Scope |
 |-----------------|-------|
 | `model_engine_creation_seconds` | Construct and load the main model engine. |
 | `draft_model_engine_creation_seconds` | Construct and load the separate draft model engine, when used. |
@@ -111,7 +111,7 @@ estimation; their final counterparts are the resources retained for serving.
 | `initial_py_executor_creation_seconds_for_kv_cache_estimation` | Construct and warm up the temporary PyExecutor used during KV cache capacity estimation. |
 | `kv_cache_capacity_configuration_seconds` | Determine final KV cache capacity from the temporary executor and available memory. |
 | `final_py_executor_creation_seconds` | Construct and warm up the final PyExecutor retained for serving. |
-| `worker_start_seconds` | Start the final executor worker. |
+| `worker_start_seconds` | Start the final PyExecutor worker. |
 | `total_executor_creation_seconds` | Complete `create_py_executor()` call, including all applicable scopes above. |
 
 The `model_loader` object contains timings for the main LLM weights. If a draft model is used,
diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py
index c79e7e1c8bcf..af70bc4ed487 100644
--- a/tensorrt_llm/executor/base_worker.py
+++ b/tensorrt_llm/executor/base_worker.py
@@ -1058,7 +1058,7 @@ def get_startup_metrics(self) -> dict:
                                                         None)
             if model_engine_metrics is not None:
                 startup_metrics[model_engine_stage] = dict(model_engine_metrics)
-        startup_metrics["executor"] = executor_metrics
+        startup_metrics["py_executor"] = executor_metrics
 
         model_engine = getattr(self.engine, "model_engine", None)
         model_loader = getattr(model_engine, "model_loader", None)
diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py
index f8e920402f68..0cb0ac895265 100644
--- a/tests/unittest/executor/test_base_worker.py
+++ b/tests/unittest/executor/test_base_worker.py
@@ -68,7 +68,7 @@ def test_get_startup_metrics_promotes_model_engine_stages():
         "final_draft_model_engine": {
             "total_warmup_seconds": 1.25
         },
-        "executor": {
+        "py_executor": {
             "worker_start_seconds": 0.5
         },
         "model_loader": {

From c2430adbc7357017c42a5d6bc6f8178c70473a94 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Mon, 24 Aug 2026 11:53:23 -0700
Subject: [PATCH 03/23] [TRTLLM-15447][fix] Address startup metrics review
 feedback

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md                  |  4 ++-
 .../_torch/pyexecutor/model_engine.py         | 33 +++++++++----------
 .../test_pytorch_model_engine_warmup.py       | 33 +++++++++++++++++++
 3 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 93751da936e0..a55c2745cab3 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -92,7 +92,8 @@ estimation creates a temporary PyExecutor before constructing the final PyExecut
 `final_model_engine` appears. The model-engine dictionaries are promoted alongside `py_executor` in
 the startup metrics payload so callers do not need to traverse another nesting level. Legacy
 two-model speculative decoding can likewise produce `initial_draft_model_engine` and
-`final_draft_model_engine`.
+`final_draft_model_engine`. When estimation is enabled, each pass performs a complete model-engine
+warmup, so both the initial and final dictionaries include the cost of a full warmup.
 
 Executor scope names describe the constructed component rather than the internal memory-category
 enum. The initial KV cache and PyExecutor scopes are temporary resources used for KV cache capacity
@@ -100,6 +101,7 @@ estimation; their final counterparts are the resources retained for serving.
 
 | PyExecutor metric | Scope |
 |-----------------|-------|
+| `config_and_checkpoint_loader_initialization_seconds` | Load the model configuration and initialize the checkpoint loader. |
 | `model_engine_creation_seconds` | Construct and load the main model engine. |
 | `draft_model_engine_creation_seconds` | Construct and load the separate draft model engine, when used. |
 | `guided_decoder_creation_seconds` | Construct guided-decoding resources. |
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 7a155559075b..8b35d1d48354 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1297,22 +1297,20 @@ def warmup_with_kv_cache_cleanup(method):
         """
 
         @functools.wraps(method)
-        def wrapper(self,
-                    resource_manager: Optional[ResourceManager] = None,
-                    *args,
-                    **kwargs):
-            result = method(self, resource_manager, *args, **kwargs)
-            with timing_metric("kv_cache_cleanup_seconds", self._metrics):
-                kv_cache_manager = (resource_manager.get_resource_manager(
-                    self.kv_cache_manager_key)
-                                    if resource_manager is not None else None)
-                if kv_cache_manager is not None:
-                    has_invalid_values = kv_cache_manager.check_invalid_values_in_kv_cache(
-                        fill_with_zero=True)
-                    if has_invalid_values:
-                        logger.warning(
-                            "NaNs/Infs have been introduced to KVCache during warmup, KVCache was filled with zeros to avoid potential issues"
-                        )
+        def wrapper(self, resource_manager: Optional[ResourceManager] = None, *args, **kwargs):
+            with timing_metric("total_warmup_seconds", self._metrics):
+                result = method(self, resource_manager, *args, **kwargs)
+                with timing_metric("kv_cache_cleanup_seconds", self._metrics):
+                    kv_cache_manager = (resource_manager.get_resource_manager(
+                        self.kv_cache_manager_key)
+                                        if resource_manager is not None else None)
+                    if kv_cache_manager is not None:
+                        has_invalid_values = kv_cache_manager.check_invalid_values_in_kv_cache(
+                            fill_with_zero=True)
+                        if has_invalid_values:
+                            logger.warning(
+                                "NaNs/Infs have been introduced to KVCache during warmup, KVCache was filled with zeros to avoid potential issues"
+                            )
             return result
 
         return wrapper
@@ -1429,8 +1427,7 @@ def forward_encoder(
     @warmup_with_kv_cache_cleanup
     def warmup(self, resource_manager: Optional[ResourceManager] = None) -> None:
         """Run model warmup and record its total wall-clock duration."""
-        with timing_metric("total_warmup_seconds", self._metrics):
-            self._warmup_impl(resource_manager)
+        self._warmup_impl(resource_manager)
 
     def _warmup_impl(self, resource_manager: Optional[ResourceManager] = None) -> None:
         """
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 0cb15b701ac8..8cc47312d476 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -1196,6 +1196,39 @@ def test_warmup_stage_metrics_are_recorded(self):
         self.assertEqual(expected_metrics, model_engine.metrics.keys())
         self.assertTrue(all(value >= 0 for value in model_engine.metrics.values()))
 
+    def test_total_warmup_metric_includes_kv_cache_cleanup(self):
+        model_engine, resource_manager = _build_engine_and_resource_manager()
+        metric_events = []
+
+        @contextlib.contextmanager
+        def record_metric_scope(metric_name, _metrics):
+            metric_events.append(("enter", metric_name))
+            try:
+                yield
+            finally:
+                metric_events.append(("exit", metric_name))
+
+        with patch(
+            "tensorrt_llm._torch.pyexecutor.model_engine.timing_metric",
+            side_effect=record_metric_scope,
+        ):
+            _run_warmup_tracked(model_engine, resource_manager)
+
+        total_and_cleanup_events = [
+            event
+            for event in metric_events
+            if event[1] in {"total_warmup_seconds", "kv_cache_cleanup_seconds"}
+        ]
+        self.assertEqual(
+            total_and_cleanup_events,
+            [
+                ("enter", "total_warmup_seconds"),
+                ("enter", "kv_cache_cleanup_seconds"),
+                ("exit", "kv_cache_cleanup_seconds"),
+                ("exit", "total_warmup_seconds"),
+            ],
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

From db427010c6973216e63c60b99d872c6e5ae6fd0c Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Mon, 24 Aug 2026 12:05:14 -0700
Subject: [PATCH 04/23] [TRTLLM-15447][docs] Document model engine warmup
 metrics

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index a55c2745cab3..8d62b51985d1 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -116,6 +116,24 @@ estimation; their final counterparts are the resources retained for serving.
 | `worker_start_seconds` | Start the final PyExecutor worker. |
 | `total_executor_creation_seconds` | Complete `create_py_executor()` call, including all applicable scopes above. |
 
+The model-engine objects contain 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 |
+|---------------------|-------|
+| `attention_warmup_seconds` | Warm up the attention backend and kernels. |
+| `general_warmup_seconds` | 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. |
+| `cuda_graph_warmup_seconds` | Run the warmup-only CUDA graph pass. |
+| `cuda_graph_capture_seconds` | Capture CUDA graphs for serving. |
+| `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` | Pre-populate the memory pool with maximum-shape allocations. |
+| `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` object contains timings for the main LLM weights. If a draft model is used,
 additional field `draft_checkpoint_preparation_seconds` and `draft_weight_population_seconds` will appear.
 A `draft_model_loader` object can also appear in deprecated 2-model style MTP setting.

From a325b570bcaeb2fff51415c19c9ad07b19d18f39 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Tue, 25 Aug 2026 20:14:47 -0700
Subject: [PATCH 05/23] misc update

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md | 31 +++++++++++++------------------
 1 file changed, 13 insertions(+), 18 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 8d62b51985d1..c6de0bf595c7 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -64,21 +64,21 @@ A typical result has the following structure:
 ```json
 {
   "py_executor": {
-    "total_executor_creation_seconds": 12.341,
-    "model_engine_creation_seconds": 5.241,
-    "final_kv_cache_creation_seconds": 0.421,
-    "final_py_executor_creation_seconds": 4.317
+    "total_executor_creation_seconds": 41.747,
+    "model_engine_creation_seconds": 1.802,
+    "final_kv_cache_creation_seconds": 0.068,
+    "final_py_executor_creation_seconds": 18.203
   },
   "final_model_engine": {
-    "total_warmup_seconds": 4.217,
-    "attention_warmup_seconds": 0.301,
-    "autotuner_warmup_seconds": 1.732,
-    "cuda_graph_capture_seconds": 1.105
+    "total_warmup_seconds": 16.885,
+    "attention_warmup_seconds": 0.884,
+    "autotuner_warmup_seconds": 0.030,
+    "cuda_graph_capture_seconds": 8.143
   },
   "model_loader": {
-    "total_model_loading_seconds":  1.971,
-    "checkpoint_preparation_seconds": 1.177,
-    "weight_population_seconds": 0.598,
+    "total_model_loading_seconds": 1.765,
+    "checkpoint_preparation_seconds": 1.001,
+    "weight_population_seconds": 0.495,
     "checkpoint_finalization_seconds": 0.031,
     "post_load_processing_seconds": 0.005
   }
@@ -92,12 +92,7 @@ estimation creates a temporary PyExecutor before constructing the final PyExecut
 `final_model_engine` appears. The model-engine dictionaries are promoted alongside `py_executor` in
 the startup metrics payload so callers do not need to traverse another nesting level. Legacy
 two-model speculative decoding can likewise produce `initial_draft_model_engine` and
-`final_draft_model_engine`. When estimation is enabled, each pass performs a complete model-engine
-warmup, so both the initial and final dictionaries include the cost of a full warmup.
-
-Executor scope names describe the constructed component rather than the internal memory-category
-enum. The initial KV cache and PyExecutor scopes are temporary resources used for KV cache capacity
-estimation; their final counterparts are the resources retained for serving.
+`final_draft_model_engine`.
 
 | PyExecutor metric | Scope |
 |-----------------|-------|
@@ -160,7 +155,7 @@ curl http://localhost:8000/server_info
 {
   "startup_metrics": {
     "model_loader": {
-      "total_model_loading_seconds": 1.971,
+      "total_model_loading_seconds": 1.765,
       ...
     }
   }

From b381bdde539dbd16333d7a488435f62a3756c970 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Tue, 25 Aug 2026 23:11:21 -0700
Subject: [PATCH 06/23] [docs] Use measured TinyLlama startup metrics

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md | 59 +++++++++++++++++++++++++++---------
 1 file changed, 44 insertions(+), 15 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index c6de0bf595c7..5c48528f9692 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -59,28 +59,57 @@ llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 print(llm.startup_metrics)
 ```
 
-A typical result has the following structure:
+A TP1 PyTorch-backend run of `TinyLlama/TinyLlama-1.1B-Chat-v1.0` produced the
+following result:
 
 ```json
 {
-  "py_executor": {
-    "total_executor_creation_seconds": 41.747,
-    "model_engine_creation_seconds": 1.802,
-    "final_kv_cache_creation_seconds": 0.068,
-    "final_py_executor_creation_seconds": 18.203
+  "initial_model_engine": {
+    "attention_warmup_seconds": 2.911,
+    "general_warmup_seconds": 0.387,
+    "autotuner_warmup_seconds": 0.041,
+    "mamba_hybrid_warmup_seconds": 0.000,
+    "cuda_graph_warmup_seconds": 0.566,
+    "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": {
-    "total_warmup_seconds": 16.885,
-    "attention_warmup_seconds": 0.884,
-    "autotuner_warmup_seconds": 0.030,
-    "cuda_graph_capture_seconds": 8.143
+    "attention_warmup_seconds": 0.057,
+    "general_warmup_seconds": 0.492,
+    "autotuner_warmup_seconds": 0.046,
+    "mamba_hybrid_warmup_seconds": 0.000,
+    "cuda_graph_warmup_seconds": 0.599,
+    "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_executor_creation_seconds": 13.373
   },
   "model_loader": {
-    "total_model_loading_seconds": 1.765,
-    "checkpoint_preparation_seconds": 1.001,
-    "weight_population_seconds": 0.495,
+    "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
   }
 }
 ```
@@ -155,7 +184,7 @@ curl http://localhost:8000/server_info
 {
   "startup_metrics": {
     "model_loader": {
-      "total_model_loading_seconds": 1.765,
+      "total_model_loading_seconds": 1.545,
       ...
     }
   }

From dd90779268aa0bd00b13932e49d8a2d55ba4574d Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 26 Aug 2026 21:34:18 -0700
Subject: [PATCH 07/23] Polished doc

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md | 47 ++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 23 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 5c48528f9692..31fd39a2db1c 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -59,8 +59,7 @@ llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
 print(llm.startup_metrics)
 ```
 
-A TP1 PyTorch-backend run of `TinyLlama/TinyLlama-1.1B-Chat-v1.0` produced the
-following result:
+A typical result has the following structure:
 
 ```json
 {
@@ -114,33 +113,34 @@ following result:
 }
 ```
 
-The `py_executor` object contains the timed scopes in `create_py_executor()`. When KV cache capacity
-estimation creates a temporary PyExecutor before constructing the final PyExecutor,
-`initial_model_engine` contains the warmup metrics from the temporary pass and
-`final_model_engine` contains the warmup metrics from the final pass. Without estimation, only
-`final_model_engine` appears. The model-engine dictionaries are promoted alongside `py_executor` in
-the startup metrics payload so callers do not need to traverse another nesting level. Legacy
-two-model speculative decoding can likewise produce `initial_draft_model_engine` and
-`final_draft_model_engine`.
+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` | Load the model configuration and initialize the checkpoint loader. |
-| `model_engine_creation_seconds` | Construct and load the main model engine. |
-| `draft_model_engine_creation_seconds` | Construct and load the separate draft model engine, when used. |
+| `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. |
-| `final_kv_cache_creation_seconds` | Construct the final KV cache retained for serving. |
-| `speculative_decoding_resource_manager_creation_seconds` | Construct speculative-decoding resource managers. |
+| `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 and warm up the temporary PyExecutor used during KV cache capacity estimation. |
+| `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_py_executor_creation_seconds` | Construct and warm up the final PyExecutor retained for serving. |
+| `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_executor_creation_seconds` | Complete `create_py_executor()` call, including all applicable scopes above. |
+| `total_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`.
 
-The model-engine objects contain the warmup scopes below. Which scopes appear depends on the model
+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.
 
@@ -158,13 +158,13 @@ stages.
 | `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` object contains timings for the main LLM weights. If a draft model is used,
-additional field `draft_checkpoint_preparation_seconds` and `draft_weight_population_seconds` will appear.
-A `draft_model_loader` object can also appear in deprecated 2-model style MTP setting.
+The `model_loader` property contains timings for loading the main LLM weights. If a draft model is used,
+additional fields `draft_checkpoint_preparation_seconds` and `draft_weight_population_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. |
@@ -172,6 +172,7 @@ A `draft_model_loader` object can also appear in deprecated 2-model style MTP se
 | `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:

From 0e6ea2e9a80714cee49b4a9105aa1f961a5d5443 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 26 Aug 2026 22:06:40 -0700
Subject: [PATCH 08/23] rename

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md                                  | 4 ++--
 tensorrt_llm/_torch/pyexecutor/py_executor_creator.py         | 2 +-
 .../executor/test_py_executor_creator_mla_cache_reuse_sync.py | 4 ++--
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 31fd39a2db1c..43d1d848772d 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -101,7 +101,7 @@ A typical result has the following structure:
     "final_kv_cache_creation_seconds": 0.057,
     "final_py_executor_creation_seconds": 4.459,
     "worker_start_seconds": 0.000,
-    "total_executor_creation_seconds": 13.373
+    "total_py_executor_creation_seconds": 13.373
   },
   "model_loader": {
     "checkpoint_preparation_seconds": 0.719,
@@ -130,7 +130,7 @@ The `py_executor` property contains the timed scopes in `create_py_executor()`,
 | `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_executor_creation_seconds` | Total time for the `create_py_executor()` call, including all applicable scopes above. |
+| `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
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index edf5b5c48b29..3e19b3bac34f 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -416,7 +416,7 @@ def _create_py_executor_impl(
         A fully initialized PyExecutor instance.
     """
     creation_metrics: dict[str, float] = {}
-    with timing_metric("total_executor_creation_seconds", creation_metrics):
+    with timing_metric("total_py_executor_creation_seconds", creation_metrics):
         py_executor = _create_py_executor(
             llm_args=llm_args,
             _startup_timer=_startup_timer,
diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
index 93793c91ea6b..b3f36e43cf9c 100644
--- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
+++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
@@ -407,7 +407,7 @@ def test_move_model_engine_metrics_moves_and_clears():
     assert draft_model_engine.metrics == {}
 
 
-def test_total_executor_creation_metric_finishes_on_error(monkeypatch):
+def test_total_py_executor_creation_metric_finishes_on_error(monkeypatch):
     captured_metrics = None
 
     @contextmanager
@@ -428,7 +428,7 @@ def _raise_during_creation(**kwargs):
     with pytest.raises(RuntimeError, match="creation failed"):
         py_executor_creator.create_py_executor(SimpleNamespace())
 
-    assert captured_metrics == {"total_executor_creation_seconds": 1.5}
+    assert captured_metrics == {"total_py_executor_creation_seconds": 1.5}
 
 
 def test_mla_unsupported_sm_fallback_syncs_cache_reuse(monkeypatch):

From dcb3a9771981b2c8fde3631ace63e1e1f9ebcf91 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 16 Sep 2026 19:05:59 -0700
Subject: [PATCH 09/23] address comments for more metrics breakdown

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md                  |  17 +-
 .../_torch/pyexecutor/model_engine.py         | 190 ++++++++++--------
 .../test_pytorch_model_engine_warmup.py       | 189 ++++++++++++++++-
 3 files changed, 300 insertions(+), 96 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 43d1d848772d..5b179cfd30bd 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -64,12 +64,13 @@ A typical result has the following structure:
 ```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,
-    "cuda_graph_warmup_seconds": 0.566,
-    "cuda_graph_capture_seconds": 0.950,
+    "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,
@@ -77,12 +78,13 @@ A typical result has the following structure:
     "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,
-    "cuda_graph_warmup_seconds": 0.599,
-    "cuda_graph_capture_seconds": 0.991,
+    "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,
@@ -146,12 +148,15 @@ stages.
 
 | Model-engine metric | Scope |
 |---------------------|-------|
+| `sampling_warmup_seconds` | Warm up sampling kernels before serving. |
 | `attention_warmup_seconds` | Warm up the attention backend and kernels. |
 | `general_warmup_seconds` | 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. |
-| `cuda_graph_warmup_seconds` | Run the warmup-only CUDA graph pass. |
-| `cuda_graph_capture_seconds` | Capture CUDA graphs for serving. |
+| `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_capture_seconds` | Capture configured prefill CUDA graphs, when applicable. |
+| `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` | Pre-populate the memory pool with maximum-shape allocations. |
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 8b35d1d48354..79194d2bc626 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1297,13 +1297,16 @@ def warmup_with_kv_cache_cleanup(method):
         """
 
         @functools.wraps(method)
-        def wrapper(self, resource_manager: Optional[ResourceManager] = None, *args, **kwargs):
+        def wrapper(self,
+                    resource_manager: Optional[ResourceManager] = None,
+                    *args,
+                    **kwargs):
             with timing_metric("total_warmup_seconds", self._metrics):
                 result = method(self, resource_manager, *args, **kwargs)
                 with timing_metric("kv_cache_cleanup_seconds", self._metrics):
                     kv_cache_manager = (resource_manager.get_resource_manager(
-                        self.kv_cache_manager_key)
-                                        if resource_manager is not None else None)
+                        self.kv_cache_manager_key) if resource_manager
+                                        is not None else None)
                     if kv_cache_manager is not None:
                         has_invalid_values = kv_cache_manager.check_invalid_values_in_kv_cache(
                             fill_with_zero=True)
@@ -1425,11 +1428,14 @@ def forward_encoder(
 
     @with_warmup_flag
     @warmup_with_kv_cache_cleanup
-    def warmup(self, resource_manager: Optional[ResourceManager] = None) -> None:
+    def warmup(self,
+               resource_manager: Optional[ResourceManager] = None) -> None:
         """Run model warmup and record its total wall-clock duration."""
         self._warmup_impl(resource_manager)
 
-    def _warmup_impl(self, resource_manager: Optional[ResourceManager] = None) -> None:
+    def _warmup_impl(self,
+                     resource_manager: Optional[ResourceManager] = None
+                     ) -> None:
         """
         Orchestrates the warmup process by calling specialized warmup methods for
         torch.compile, the autotuner, and CUDA graphs.
@@ -1458,7 +1464,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
         # cuda_graph_config=None flashinfer's sampling kernels would be
         # JIT-built mid-serving.
         self._eager_workspace_reclaimer = None
-        with self._warmup_timer.phase("sampling_module_prewarm"):
+        with self._warmup_timer.phase("sampling_module_prewarm"), timing_metric(
+                "sampling_warmup_seconds", self._metrics):
             warmup_sampling_module()
             if self.enable_in_graph_sampling:
                 # The fast tier samples inside the captured graph via a
@@ -1562,15 +1569,11 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             with self.cuda_graph_runner.allow_capture():
                 self.cuda_graph_runner.is_warmup_only = True
                 try:
-                    with timing_metric("cuda_graph_warmup_seconds",
-                                       self._metrics):
-                        with self.maybe_autotune_lora():
-                            self._run_cuda_graph_warmup(resource_manager)
+                    self._run_cuda_graph_warmup(resource_manager)
                 finally:
                     self.cuda_graph_runner.is_warmup_only = False
                 self.cuda_graph_runner.padding_dummy_requests = {}
-                with timing_metric("cuda_graph_capture_seconds", self._metrics):
-                    self._run_cuda_graph_warmup(resource_manager)
+                self._run_cuda_graph_warmup(resource_manager)
         log_mem_snapshot("warmup/after_cuda_graph_capture")
         # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned
         # batch bucket the runtime can produce (max_batch_size scaled by the
@@ -2613,34 +2616,44 @@ def _get_graphs_to_capture(
 
     def _run_cuda_graph_warmup(self, resource_manager: ResourceManager):
         """Warm up or capture CUDA graphs for the configured graph shapes."""
-        if not (self.cuda_graph_runner.enabled
-                or self.prefill_cuda_graph_backend
-                != PrefillCudaGraphBackend.DISABLED):
-            return
-
-        from ..modules.linear import (MXFP8LinearMethod,
-                                      flashinfer_mxfp8_autotune,
-                                      flashinfer_mxfp8_decode_graph_capture)
-
-        # The automatic MiniMax-M3 MXFP8 selection is decode-graph-only.
-        # Tune every generation graph shape during the warmup-only pass. Keep
-        # piecewise context/prefill graph capture on the native backend.
-        flashinfer_methods = [
-            quant_method for module in self.model.modules()
-            if isinstance((quant_method := getattr(module, "quant_method", None)
-                           ), MXFP8LinearMethod)
-            and quant_method.needs_flashinfer_autotune
-        ]
-        flashinfer_autotune_context = (
-            flashinfer_mxfp8_autotune() if self.cuda_graph_runner.is_warmup_only
-            and flashinfer_methods else contextlib.nullcontext())
-        with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture(
-        ):
-            self._capture_generation_cuda_graphs(resource_manager)
-        self._capture_mixed_encoder_decoder_cuda_graphs(resource_manager)
+        is_warmup_only = self.cuda_graph_runner.is_warmup_only
+        metric_name = ("gen_cuda_graph_warmup_seconds"
+                       if is_warmup_only else "gen_cuda_graph_capture_seconds")
+        with timing_metric(metric_name, self._metrics):
+            # Include LoRA autotuning and its PP cache hand-off/cleanup in
+            # warmup timing, including when this rank has no graph shapes.
+            lora_context = (self.maybe_autotune_lora()
+                            if is_warmup_only else contextlib.nullcontext())
+            with lora_context:
+                if not (self.cuda_graph_runner.enabled
+                        or self.prefill_cuda_graph_backend
+                        != PrefillCudaGraphBackend.DISABLED):
+                    return
+
+                from ..modules.linear import (
+                    MXFP8LinearMethod, flashinfer_mxfp8_autotune,
+                    flashinfer_mxfp8_decode_graph_capture)
+
+                # The automatic MiniMax-M3 MXFP8 selection is decode-graph-only.
+                # Tune every generation graph shape during the warmup-only pass.
+                # Keep piecewise context/prefill capture on the native backend.
+                flashinfer_methods = [
+                    quant_method for module in self.model.modules()
+                    if isinstance((quant_method := getattr(
+                        module, "quant_method", None)), MXFP8LinearMethod)
+                    and quant_method.needs_flashinfer_autotune
+                ]
+                flashinfer_autotune_context = (
+                    flashinfer_mxfp8_autotune() if is_warmup_only
+                    and flashinfer_methods else contextlib.nullcontext())
+                with flashinfer_autotune_context, flashinfer_mxfp8_decode_graph_capture(
+                ):
+                    self._capture_generation_cuda_graphs(resource_manager)
+                self._capture_mixed_encoder_decoder_cuda_graphs(
+                    resource_manager)
         # Piecewise graphs have separate capture machinery and do not use the
         # whole-model attention workspace. Capture them only on the second pass.
-        if not self.cuda_graph_runner.is_warmup_only:
+        if not is_warmup_only:
             self._capture_prefill_cuda_graphs(resource_manager)
 
     def _capture_generation_cuda_graphs(self,
@@ -3021,7 +3034,12 @@ def _capture_mixed_encoder_decoder_cuda_graphs(
                         self.runtime_draft_len = saved_runtime_draft_len
 
     def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager):
-        """Capture configured CUDA graphs for context/prefill steps."""
+        """Capture prefill CUDA graphs, then warm up logits buffer allocations.
+
+        After capture, run prefill batches with many requests to warm up
+        logits buffers outside the captured model body. This post-capture
+        warmup runs for both piecewise and breakable prefill CUDA graphs.
+        """
         if (self.prefill_cuda_graph_backend == PrefillCudaGraphBackend.DISABLED
                 or (self.prefill_cuda_graph_backend
                     == PrefillCudaGraphBackend.PIECEWISE
@@ -3035,62 +3053,68 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager):
         capture_context = (capture_piecewise_cuda_graph(True)
                            if self._torch_compile_piecewise_cuda_graph else
                            contextlib.nullcontext())
-        with capture_context, self.no_cuda_graph():
+        with timing_metric("ctx_cuda_graph_capture_seconds", self._metrics):
+            with capture_context, self.no_cuda_graph():
+                for num_tokens in prefill_cuda_graph_num_tokens:
+                    warmup_request = self._create_warmup_request(
+                        resource_manager, num_tokens, 0)
+                    with self._release_batch_context(warmup_request,
+                                                     resource_manager) as batch:
+                        self._assert_all_tp_ranks_have_warmup_batch(
+                            batch, num_tokens)
+                        if batch is None:
+                            continue
+
+                        logger.info(
+                            f"Run prefill CUDA graph capture for num tokens={num_tokens}"
+                        )
+                        if self.breakable_cuda_graph_runner is not None:
+                            self.breakable_cuda_graph_runner.capture(
+                                num_tokens, lambda: self.forward(
+                                    batch,
+                                    new_tensors_device=None,
+                                    resource_manager=resource_manager))
+                        else:
+                            # Run a few times to ensure torch.compile capture.
+                            for _ in range(4):
+                                self.forward(batch,
+                                             new_tensors_device=None,
+                                             resource_manager=resource_manager)
+
+                        torch.cuda.synchronize()
+                        gc.collect()
+                        torch.cuda.empty_cache()
+
+        # The logits allocations grow with the number of requests and are not
+        # part of the captured model body. Warm up the largest request count so
+        # those allocations can be reused during stable inference.
+        with timing_metric("post_ctx_cuda_graph_capture_warmup_seconds",
+                           self._metrics):
             for num_tokens in prefill_cuda_graph_num_tokens:
                 warmup_request = self._create_warmup_request(
-                    resource_manager, num_tokens, 0)
+                    resource_manager, num_tokens, 0, least_requests=False)
                 with self._release_batch_context(warmup_request,
                                                  resource_manager) as batch:
                     self._assert_all_tp_ranks_have_warmup_batch(
                         batch, num_tokens)
                     if batch is None:
                         continue
-
                     logger.info(
-                        f"Run prefill CUDA graph capture for num tokens={num_tokens}"
+                        f"Run prefill CUDA graph warmup for num tokens={num_tokens} with most requests"
                     )
                     if self.breakable_cuda_graph_runner is not None:
-                        self.breakable_cuda_graph_runner.capture(
-                            num_tokens, lambda: self.forward(
-                                batch,
-                                new_tensors_device=None,
-                                resource_manager=resource_manager))
+                        with self.no_cuda_graph():
+                            self.breakable_cuda_graph_runner.warmup(
+                                lambda: self.forward(batch,
+                                                     new_tensors_device=None,
+                                                     resource_manager=
+                                                     resource_manager),
+                                steps=1)
                     else:
-                        # Run a few times to ensure torch.compile capture.
-                        for _ in range(4):
-                            self.forward(batch,
-                                         new_tensors_device=None,
-                                         resource_manager=resource_manager)
-
-        # The logits allocations grow with the number of requests and are not
-        # part of the captured model body. Warm up the largest request count so
-        # those allocations can be reused during stable inference.
-        for num_tokens in prefill_cuda_graph_num_tokens:
-            warmup_request = self._create_warmup_request(resource_manager,
-                                                         num_tokens,
-                                                         0,
-                                                         least_requests=False)
-            with self._release_batch_context(warmup_request,
-                                             resource_manager) as batch:
-                self._assert_all_tp_ranks_have_warmup_batch(batch, num_tokens)
-                if batch is None:
-                    continue
-                logger.info(
-                    f"Run prefill CUDA graph warmup for num tokens={num_tokens} with most requests"
-                )
-                if self.breakable_cuda_graph_runner is not None:
-                    with self.no_cuda_graph():
-                        self.breakable_cuda_graph_runner.warmup(
-                            lambda: self.forward(batch,
-                                                 new_tensors_device=None,
-                                                 resource_manager=
-                                                 resource_manager),
-                            steps=1)
-                else:
-                    self.forward(batch,
-                                 new_tensors_device=None,
-                                 resource_manager=resource_manager)
-                torch.cuda.synchronize()
+                        self.forward(batch,
+                                     new_tensors_device=None,
+                                     resource_manager=resource_manager)
+                    torch.cuda.synchronize()
 
     ### Helper methods promoted from the original warmup method ###
 
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 8cc47312d476..6a1f2ab3b9e6 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -15,6 +15,7 @@
 import sys
 import unittest
 from collections import OrderedDict
+from collections.abc import Iterator
 from dataclasses import dataclass
 from types import ModuleType, SimpleNamespace
 from unittest.mock import Mock, call, patch
@@ -601,9 +602,11 @@ def _wrapped(*args, **kwargs):
 def _run_warmup_tracked(
     model_engine, resource_manager, *, force_helix_cp=False, capture_logs=False
 ):
-    """Patch the four warmup helpers + empty_cache + MoERunner.clear and run
-    model_engine.warmup(). Optionally force helix CP and capture logs.
-    Returns (call_order_list, log_records_or_None)."""
+    """Patch warmup stages and cleanup operations, then run warmup.
+
+    Optionally force helix CP and capture logs. Returns
+    (call_order_list, log_records_or_None).
+    """
     tracker = _Tracker()
     helix_ctx = (
         patch.object(model_engine.mapping, "has_cp_helix", return_value=True)
@@ -615,7 +618,25 @@ def _run_warmup_tracked(
         helix_ctx,
         patch.object(model_engine, "_general_warmup", side_effect=tracker("general_warmup")),
         patch.object(model_engine, "_run_autotuner_warmup", side_effect=tracker("autotuner")),
-        patch.object(model_engine, "_run_cuda_graph_warmup", side_effect=tracker("cuda_graph")),
+        patch.object(
+            model_engine,
+            "_capture_generation_cuda_graphs",
+            side_effect=tracker("generation_cuda_graph"),
+        ),
+        patch.object(
+            model_engine,
+            "_capture_mixed_encoder_decoder_cuda_graphs",
+            side_effect=tracker("mixed_cuda_graph"),
+        ),
+        patch.object(
+            model_engine,
+            "_capture_prefill_cuda_graphs",
+            side_effect=tracker("prefill_cuda_graph"),
+        ),
+        patch(
+            "tensorrt_llm._torch.pyexecutor.model_engine.warmup_sampling_module",
+            side_effect=tracker("sampling_warmup"),
+        ),
         patch("torch.cuda.empty_cache", side_effect=tracker("empty_cache")),
         patch(
             "tensorrt_llm._torch.custom_ops.torch_custom_ops.MoERunner.clear_all_workspaces",
@@ -658,6 +679,7 @@ def test_no_kv_cache_warmup_delegates_runner_lifecycle(self):
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))
         model_engine.moe_load_balancer = None
+        model_engine._metrics = {}
         model_engine.is_warmup = False
         model_engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
         model_engine._fallback_to_engine = False
@@ -678,6 +700,7 @@ def test_no_kv_cache_warmup_delegates_runner_lifecycle(self):
 
     def test_no_kv_cache_warmup_rejects_allocated_kv_cache(self):
         model_engine = object.__new__(PyTorchModelEngine)
+        model_engine._metrics = {}
         model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))
         model_engine.moe_load_balancer = None
         model_engine.is_warmup = False
@@ -702,6 +725,7 @@ def test_no_kv_cache_warmup_rejects_allocated_kv_cache(self):
     def test_legacy_warmup_skips_without_kv_cache(self):
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine._warmup_timer = _WarmupTimer(rank=0)
+        model_engine._metrics = {}
         model_engine.moe_load_balancer = None
         model_engine.is_warmup = False
         model_engine.enable_in_graph_sampling = False
@@ -739,6 +763,157 @@ def test_encoder_decoder_encoder_warmup_delegates_runner_lifecycle(self):
             [call.warmup(resource_manager)],
         )
 
+    def test_cuda_graph_metrics_exclude_piecewise_stages(self) -> None:
+        model_engine = object.__new__(PyTorchModelEngine)
+        model_engine.model = SimpleNamespace(modules=lambda: [])
+        model_engine.llm_args = SimpleNamespace(enable_autotuner=True)
+        model_engine.cuda_graph_lora_manager = object()
+        model_engine.cuda_graph_runner = SimpleNamespace(
+            enabled=True,
+            is_warmup_only=False,
+        )
+        model_engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE
+        model_engine._metrics = {}
+        resource_manager = object()
+        events = []
+
+        @contextlib.contextmanager
+        def record_metric_scope(metric_name: str, _metrics: dict[str, float]) -> Iterator[None]:
+            events.append(("enter", metric_name))
+            try:
+                yield
+            finally:
+                events.append(("exit", metric_name))
+
+        autotuner = SimpleNamespace(
+            cache_pp_recv=lambda: events.append(("lora", "cache_pp_recv")),
+            cache_pp_send=lambda: events.append(("lora", "cache_pp_send")),
+            clean_pp_flag=lambda: events.append(("lora", "clean_pp_flag")),
+        )
+        lora_cleanup_events = [
+            ("lora", "cache_pp_recv"),
+            ("lora", "cache_pp_send"),
+            ("lora", "clean_pp_flag"),
+            ("exit", "lora_autotune"),
+            ("exit", "gen_cuda_graph_warmup_seconds"),
+        ]
+
+        with (
+            patch(
+                "tensorrt_llm._torch.pyexecutor.model_engine.autotune",
+                side_effect=lambda **_: record_metric_scope("lora_autotune", {}),
+            ),
+            patch(
+                "tensorrt_llm._torch.pyexecutor.model_engine.AutoTuner.get",
+                return_value=autotuner,
+            ),
+            patch(
+                "tensorrt_llm._torch.modules.linear.flashinfer_mxfp8_decode_graph_capture",
+                side_effect=contextlib.nullcontext,
+            ),
+            patch(
+                "tensorrt_llm._torch.pyexecutor.model_engine.timing_metric",
+                side_effect=record_metric_scope,
+            ),
+            patch.object(
+                model_engine,
+                "_capture_generation_cuda_graphs",
+                side_effect=lambda _: events.append(("stage", "generation")),
+            ) as generation,
+            patch.object(
+                model_engine,
+                "_capture_mixed_encoder_decoder_cuda_graphs",
+                side_effect=lambda _: events.append(("stage", "mixed")),
+            ),
+            patch.object(
+                model_engine,
+                "_capture_prefill_cuda_graphs",
+                side_effect=lambda _: events.append(("stage", "piecewise")),
+            ) as piecewise,
+        ):
+            model_engine._run_cuda_graph_warmup(resource_manager)
+
+            self.assertEqual(
+                events,
+                [
+                    ("enter", "gen_cuda_graph_capture_seconds"),
+                    ("stage", "generation"),
+                    ("stage", "mixed"),
+                    ("exit", "gen_cuda_graph_capture_seconds"),
+                    ("stage", "piecewise"),
+                ],
+            )
+
+            events.clear()
+            piecewise.reset_mock()
+            model_engine.cuda_graph_runner.is_warmup_only = True
+            model_engine._run_cuda_graph_warmup(resource_manager)
+
+            self.assertEqual(
+                events,
+                [
+                    ("enter", "gen_cuda_graph_warmup_seconds"),
+                    ("enter", "lora_autotune"),
+                    ("stage", "generation"),
+                    ("stage", "mixed"),
+                    *lora_cleanup_events,
+                ],
+            )
+            piecewise.assert_not_called()
+
+            events.clear()
+            generation.side_effect = RuntimeError("warmup failed")
+            with self.assertRaisesRegex(RuntimeError, "warmup failed"):
+                model_engine._run_cuda_graph_warmup(resource_manager)
+            self.assertEqual(
+                events,
+                [
+                    ("enter", "gen_cuda_graph_warmup_seconds"),
+                    ("enter", "lora_autotune"),
+                    *lora_cleanup_events,
+                ],
+            )
+            piecewise.assert_not_called()
+
+            events.clear()
+            generation.reset_mock()
+            model_engine.cuda_graph_runner.enabled = False
+            model_engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.DISABLED
+            model_engine._run_cuda_graph_warmup(resource_manager)
+            self.assertEqual(
+                events,
+                [
+                    ("enter", "gen_cuda_graph_warmup_seconds"),
+                    ("enter", "lora_autotune"),
+                    *lora_cleanup_events,
+                ],
+            )
+            generation.assert_not_called()
+            piecewise.assert_not_called()
+    def test_piecewise_cuda_graph_metrics_are_recorded_separately(self) -> None:
+        model_engine = object.__new__(PyTorchModelEngine)
+        model_engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE
+        model_engine._torch_compile_enabled = True
+        model_engine._torch_compile_piecewise_cuda_graph = False
+        model_engine._prefill_cuda_graph_num_tokens = []
+        model_engine._metrics = {}
+
+        with patch.object(
+            model_engine,
+            "no_cuda_graph",
+            return_value=contextlib.nullcontext(),
+        ):
+            model_engine._capture_prefill_cuda_graphs(object())
+
+        self.assertEqual(
+            model_engine.metrics.keys(),
+            {
+                "ctx_cuda_graph_capture_seconds",
+                "post_ctx_cuda_graph_capture_warmup_seconds",
+            },
+        )
+        self.assertTrue(all(value >= 0 for value in model_engine.metrics.values()))
+
     def test_empty_cache_fires_immediately_after_autotuner(self):
         """Change 1 placement: empty_cache must be the call right after
         _run_autotuner_warmup."""
@@ -1181,12 +1356,13 @@ def test_warmup_stage_metrics_are_recorded(self):
         _run_warmup_tracked(model_engine, resource_manager)
 
         expected_metrics = {
+            "sampling_warmup_seconds",
             "attention_warmup_seconds",
             "general_warmup_seconds",
             "autotuner_warmup_seconds",
             "mamba_hybrid_warmup_seconds",
-            "cuda_graph_warmup_seconds",
-            "cuda_graph_capture_seconds",
+            "gen_cuda_graph_warmup_seconds",
+            "gen_cuda_graph_capture_seconds",
             "dg_paged_mqa_warmup_seconds",
             "cute_dsl_radix_topk_warmup_seconds",
             "memory_pool_prepopulation_seconds",
@@ -1229,6 +1405,5 @@ def record_metric_scope(metric_name, _metrics):
             ],
         )
 
-
 if __name__ == "__main__":
     unittest.main()

From f81ccc101c016d6cc4a5cbfb8e8a7696111aaae8 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 16 Sep 2026 19:18:12 -0700
Subject: [PATCH 10/23] clean up tests

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../test_pytorch_model_engine_warmup.py       | 57 +------------------
 tests/unittest/executor/test_base_worker.py   |  1 +
 2 files changed, 4 insertions(+), 54 deletions(-)

diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 6a1f2ab3b9e6..eadf711afcdd 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -763,6 +763,7 @@ def test_encoder_decoder_encoder_warmup_delegates_runner_lifecycle(self):
             [call.warmup(resource_manager)],
         )
 
+    @pytest.mark.cpu_only
     def test_cuda_graph_metrics_exclude_piecewise_stages(self) -> None:
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.model = SimpleNamespace(modules=lambda: [])
@@ -890,6 +891,8 @@ def record_metric_scope(metric_name: str, _metrics: dict[str, float]) -> Iterato
             )
             generation.assert_not_called()
             piecewise.assert_not_called()
+
+    @pytest.mark.cpu_only
     def test_piecewise_cuda_graph_metrics_are_recorded_separately(self) -> None:
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE
@@ -1351,59 +1354,5 @@ def test_native_mxfp8_respects_disabled_global_autotuner(self):
         self.assertFalse(method.use_native_autotuner)
         self.assertFalse(method.needs_native_autotune)
 
-    def test_warmup_stage_metrics_are_recorded(self):
-        model_engine, resource_manager = _build_engine_and_resource_manager()
-        _run_warmup_tracked(model_engine, resource_manager)
-
-        expected_metrics = {
-            "sampling_warmup_seconds",
-            "attention_warmup_seconds",
-            "general_warmup_seconds",
-            "autotuner_warmup_seconds",
-            "mamba_hybrid_warmup_seconds",
-            "gen_cuda_graph_warmup_seconds",
-            "gen_cuda_graph_capture_seconds",
-            "dg_paged_mqa_warmup_seconds",
-            "cute_dsl_radix_topk_warmup_seconds",
-            "memory_pool_prepopulation_seconds",
-            "kv_cache_cleanup_seconds",
-            "total_warmup_seconds",
-        }
-        self.assertEqual(expected_metrics, model_engine.metrics.keys())
-        self.assertTrue(all(value >= 0 for value in model_engine.metrics.values()))
-
-    def test_total_warmup_metric_includes_kv_cache_cleanup(self):
-        model_engine, resource_manager = _build_engine_and_resource_manager()
-        metric_events = []
-
-        @contextlib.contextmanager
-        def record_metric_scope(metric_name, _metrics):
-            metric_events.append(("enter", metric_name))
-            try:
-                yield
-            finally:
-                metric_events.append(("exit", metric_name))
-
-        with patch(
-            "tensorrt_llm._torch.pyexecutor.model_engine.timing_metric",
-            side_effect=record_metric_scope,
-        ):
-            _run_warmup_tracked(model_engine, resource_manager)
-
-        total_and_cleanup_events = [
-            event
-            for event in metric_events
-            if event[1] in {"total_warmup_seconds", "kv_cache_cleanup_seconds"}
-        ]
-        self.assertEqual(
-            total_and_cleanup_events,
-            [
-                ("enter", "total_warmup_seconds"),
-                ("enter", "kv_cache_cleanup_seconds"),
-                ("exit", "kv_cache_cleanup_seconds"),
-                ("exit", "total_warmup_seconds"),
-            ],
-        )
-
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py
index 0cb0ac895265..e9879edde0a6 100644
--- a/tests/unittest/executor/test_base_worker.py
+++ b/tests/unittest/executor/test_base_worker.py
@@ -24,6 +24,7 @@
 model_path = llm_models_root() / default_model_name
 
 
+@pytest.mark.cpu_only
 def test_get_startup_metrics_promotes_model_engine_stages():
     worker = object.__new__(BaseWorker)
     worker._is_pytorch_backend = True

From 08e7a2cdd3891e7dfc8528dd57f730ef71c3b919 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 16 Sep 2026 19:28:02 -0700
Subject: [PATCH 11/23] Address startup metrics documentation and test review
 comments

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md                           | 10 ++++++----
 .../executor/test_pytorch_model_engine_warmup.py       |  1 +
 tests/unittest/executor/test_base_worker.py            |  2 +-
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 5b179cfd30bd..50fde01ec5ab 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -59,7 +59,7 @@ 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
 {
@@ -163,8 +163,10 @@ stages.
 | `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. If a draft model is used,
-additional fields `draft_checkpoint_preparation_seconds` and `draft_weight_population_seconds` will appear.
+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.
 
@@ -172,7 +174,7 @@ The `py_executor.model_engine_creation_seconds` timing includes the total `model
 |--------|-------------|
 | `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. |
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index eadf711afcdd..6b82e766c752 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -1354,5 +1354,6 @@ def test_native_mxfp8_respects_disabled_global_autotuner(self):
         self.assertFalse(method.use_native_autotuner)
         self.assertFalse(method.needs_native_autotune)
 
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py
index e9879edde0a6..de24dc6fd0ae 100644
--- a/tests/unittest/executor/test_base_worker.py
+++ b/tests/unittest/executor/test_base_worker.py
@@ -25,7 +25,7 @@
 
 
 @pytest.mark.cpu_only
-def test_get_startup_metrics_promotes_model_engine_stages():
+def test_get_startup_metrics_promotes_model_engine_stages() -> None:
     worker = object.__new__(BaseWorker)
     worker._is_pytorch_backend = True
     worker.engine = SimpleNamespace(

From 320b61deea62564e1c108dd9b41ff718ec95e889 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 16 Sep 2026 20:24:08 -0700
Subject: [PATCH 12/23] Assert metrics reset after initial startup-stage
 transfer

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../executor/test_py_executor_creator_mla_cache_reuse_sync.py  | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
index b3f36e43cf9c..0eed05c0cafd 100644
--- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
+++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
@@ -391,6 +391,9 @@ def test_move_model_engine_metrics_moves_and_clears():
     py_executor_creator._move_model_engine_metrics(
         py_executor, model_engine, "initial", draft_model_engine
     )
+    assert model_engine.metrics == {}
+    assert draft_model_engine.metrics == {}
+
     model_engine.metrics["total_warmup_seconds"] = 3.5
     draft_model_engine.metrics["total_warmup_seconds"] = 2.0
     py_executor_creator._move_model_engine_metrics(

From 1037880e11a76c00a0c8d5cb5a22be4a3727c0e2 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Tue, 22 Sep 2026 15:36:19 -0700
Subject: [PATCH 13/23] fix sampling timing metric scope and add prefill CUDA
 graph warmup timing metric

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md                  |   5 +-
 .../_torch/compilation/piecewise_optimizer.py |   7 +-
 .../pyexecutor/breakable_cuda_graph_runner.py |  51 +++--
 .../_torch/pyexecutor/model_engine.py         |  73 ++++---
 .../test_pytorch_model_engine_warmup.py       | 192 +++++++++++++++---
 5 files changed, 258 insertions(+), 70 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index 50fde01ec5ab..d6fa516bf1e4 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -148,14 +148,15 @@ stages.
 
 | Model-engine metric | Scope |
 |---------------------|-------|
-| `sampling_warmup_seconds` | Warm up sampling kernels before serving. |
+| `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` | 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_capture_seconds` | Capture configured prefill CUDA graphs, when applicable. |
+| `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. |
diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py
index 5ade1813c7c1..2a964aa0a209 100644
--- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py
+++ b/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)
 
diff --git a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py
index 9eccb529196d..8f683a208da0 100644
--- a/tensorrt_llm/_torch/pyexecutor/breakable_cuda_graph_runner.py
+++ b/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):
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 79194d2bc626..6d8e1391ecff 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -8,6 +8,7 @@
 import math
 import os
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from contextlib import contextmanager
 from typing import (Any, Callable, Dict, Iterator, List, Optional, Sequence,
                     Tuple, Type, Union)
@@ -47,6 +48,7 @@
 from ..attention.backends.utils import get_attention_backend
 from ..autotuner import AutoTuner, autotune
 from ..compilation.backend import Backend
+from ..compilation.piecewise_optimizer import PiecewiseRunner
 from ..compilation.utils import capture_piecewise_cuda_graph
 from ..distributed import Distributed
 from ..distributed.communicator import init_pp_comm
@@ -378,7 +380,7 @@ def __init__(
     ):
         _configure_deep_gemm_pdl()
 
-        self._metrics: dict[str, float] = {}
+        self._metrics: dict[str, float] = defaultdict(float)
         self.forward_pass_callable = None
         self._cleanup_done = False
         self._model_caller: Optional[ModelCaller] = None
@@ -3034,7 +3036,7 @@ def _capture_mixed_encoder_decoder_cuda_graphs(
                         self.runtime_draft_len = saved_runtime_draft_len
 
     def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager):
-        """Capture prefill CUDA graphs, then warm up logits buffer allocations.
+        """Warm up and capture prefill CUDA graphs, timing each phase separately.
 
         After capture, run prefill batches with many requests to warm up
         logits buffers outside the captured model body. This post-capture
@@ -3053,37 +3055,52 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager):
         capture_context = (capture_piecewise_cuda_graph(True)
                            if self._torch_compile_piecewise_cuda_graph else
                            contextlib.nullcontext())
-        with timing_metric("ctx_cuda_graph_capture_seconds", self._metrics):
-            with capture_context, self.no_cuda_graph():
-                for num_tokens in prefill_cuda_graph_num_tokens:
-                    warmup_request = self._create_warmup_request(
-                        resource_manager, num_tokens, 0)
-                    with self._release_batch_context(warmup_request,
-                                                     resource_manager) as batch:
-                        self._assert_all_tp_ranks_have_warmup_batch(
-                            batch, num_tokens)
-                        if batch is None:
-                            continue
+        with capture_context, self.no_cuda_graph():
+            for num_tokens in prefill_cuda_graph_num_tokens:
+                warmup_request = self._create_warmup_request(
+                    resource_manager, num_tokens, 0)
+                with self._release_batch_context(warmup_request,
+                                                 resource_manager) as batch:
+                    self._assert_all_tp_ranks_have_warmup_batch(
+                        batch, num_tokens)
+                    if batch is None:
+                        continue
 
-                        logger.info(
-                            f"Run prefill CUDA graph capture for num tokens={num_tokens}"
-                        )
-                        if self.breakable_cuda_graph_runner is not None:
-                            self.breakable_cuda_graph_runner.capture(
-                                num_tokens, lambda: self.forward(
-                                    batch,
-                                    new_tensors_device=None,
-                                    resource_manager=resource_manager))
-                        else:
-                            # Run a few times to ensure torch.compile capture.
-                            for _ in range(4):
+                    logger.info(
+                        f"Run prefill CUDA graph capture for num tokens={num_tokens}"
+                    )
+                    if self.breakable_cuda_graph_runner is not None:
+                        runner = self.breakable_cuda_graph_runner
+                        runner.capture(
+                            num_tokens, lambda: self.forward(
+                                batch,
+                                new_tensors_device=None,
+                                resource_manager=resource_manager))
+                        self._metrics[
+                            "ctx_cuda_graph_warmup_seconds"] += runner.metrics.get(
+                                BreakableCUDAGraphRunner.
+                                CUDA_GRAPH_WARMUP_METRIC, 0.0)
+                        self._metrics[
+                            "ctx_cuda_graph_capture_seconds"] += runner.metrics.get(
+                                BreakableCUDAGraphRunner.
+                                CUDA_GRAPH_CAPTURE_METRIC, 0.0)
+                    else:
+                        with timing_metric("ctx_cuda_graph_warmup_seconds",
+                                           self._metrics):
+                            for _ in range(PiecewiseRunner.WARMUP_STEPS):
                                 self.forward(batch,
                                              new_tensors_device=None,
                                              resource_manager=resource_manager)
+                            torch.cuda.synchronize()
+                        with timing_metric("ctx_cuda_graph_capture_seconds",
+                                           self._metrics):
+                            self.forward(batch,
+                                         new_tensors_device=None,
+                                         resource_manager=resource_manager)
+                            torch.cuda.synchronize()
 
-                        torch.cuda.synchronize()
-                        gc.collect()
-                        torch.cuda.empty_cache()
+                    gc.collect()
+                    torch.cuda.empty_cache()
 
         # The logits allocations grow with the number of requests and are not
         # part of the captured model body. Warm up the largest request count so
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 6b82e766c752..c438915390d4 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -14,7 +14,7 @@
 import os
 import sys
 import unittest
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from collections.abc import Iterator
 from dataclasses import dataclass
 from types import ModuleType, SimpleNamespace
@@ -29,8 +29,9 @@
 from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8GemmRunner
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
-from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM
+from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM, timing_metric
 from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod
+from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner
 from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller
 from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner
 from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner
@@ -722,7 +723,7 @@ def test_no_kv_cache_warmup_rejects_allocated_kv_cache(self):
         model_engine._runner.warmup.assert_called_once_with(resource_manager)
 
     @pytest.mark.cpu_only
-    def test_legacy_warmup_skips_without_kv_cache(self):
+    def test_legacy_warmup_skips_without_kv_cache(self) -> None:
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine._warmup_timer = _WarmupTimer(rank=0)
         model_engine._metrics = {}
@@ -732,21 +733,62 @@ def test_legacy_warmup_skips_without_kv_cache(self):
         model_engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
         model_engine._fallback_to_engine = True
         model_engine._runner = None
+        model_engine.model = SimpleNamespace(config=SimpleNamespace(vocab_size=128))
+        model_engine.dtype = torch.float16
+        model_engine._cuda_graph_batch_sizes = [1, 4]
         resource_manager = Mock()
         resource_manager.get_resource_manager.return_value = None
+        events = []
+
+        @contextlib.contextmanager
+        def record_metric_scope(name: str, metrics: dict[str, float]) -> Iterator[None]:
+            events.append(("enter", name))
+            with timing_metric(name, metrics):
+                yield
+            events.append(("exit", name))
 
         with (
             patch(
-                "tensorrt_llm._torch.pyexecutor.model_engine.warmup_sampling_module"
+                "tensorrt_llm._torch.pyexecutor.model_engine.timing_metric",
+                side_effect=record_metric_scope,
+            ),
+            patch(
+                "tensorrt_llm._torch.pyexecutor.model_engine.warmup_sampling_module",
+                side_effect=lambda: events.append("sampling"),
             ) as warmup_sampling,
+            patch(
+                "tensorrt_llm._torch.pyexecutor.model_engine.warmup_sample_from_logits_op",
+                side_effect=lambda *args: events.append("in_graph_sampling"),
+            ) as warmup_in_graph_sampling,
             _capture_tllm_logs() as logs,
         ):
-            model_engine.warmup(resource_manager)
+            for enabled in (False, True):
+                with self.subTest(enable_in_graph_sampling=enabled):
+                    events.clear()
+                    warmup_sampling.reset_mock()
+                    warmup_in_graph_sampling.reset_mock()
+                    model_engine.enable_in_graph_sampling = enabled
+                    model_engine._eager_workspace_reclaimer = object()
+                    model_engine._warmup_impl(resource_manager)
+                    self.assertEqual(
+                        events,
+                        [("enter", "sampling_warmup_seconds"), "sampling"]
+                        + (["in_graph_sampling"] if enabled else [])
+                        + [("exit", "sampling_warmup_seconds")],
+                    )
+                    self.assertIsNone(model_engine._eager_workspace_reclaimer)
+                    warmup_sampling.assert_called_once_with()
+                    if enabled:
+                        warmup_in_graph_sampling.assert_called_once_with(
+                            128, torch.device("cuda"), torch.float16, [1, 4]
+                        )
+                    else:
+                        warmup_in_graph_sampling.assert_not_called()
 
         self.assertTrue(
             any("Skipping warm up as no KV Cache manager allocated." in log for log in logs)
         )
-        warmup_sampling.assert_called_once_with()
+        self.assertGreaterEqual(model_engine.metrics["sampling_warmup_seconds"], 0)
 
     def test_encoder_decoder_encoder_warmup_delegates_runner_lifecycle(self):
         model_engine = object.__new__(PyTorchModelEngine)
@@ -897,25 +939,129 @@ def test_piecewise_cuda_graph_metrics_are_recorded_separately(self) -> None:
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE
         model_engine._torch_compile_enabled = True
-        model_engine._torch_compile_piecewise_cuda_graph = False
-        model_engine._prefill_cuda_graph_num_tokens = []
-        model_engine._metrics = {}
+        model_engine._torch_compile_piecewise_cuda_graph = True
+        model_engine._prefill_cuda_graph_num_tokens = [4, 8]
+        model_engine._metrics = defaultdict(float)
+        resource_manager = object()
+        events = []
+        scopes = []
+        warmup_metric = "ctx_cuda_graph_warmup_seconds"
+        capture_metric = "ctx_cuda_graph_capture_seconds"
+        post_metric = "post_ctx_cuda_graph_capture_warmup_seconds"
 
-        with patch.object(
-            model_engine,
-            "no_cuda_graph",
-            return_value=contextlib.nullcontext(),
-        ):
-            model_engine._capture_prefill_cuda_graphs(object())
+        @contextlib.contextmanager
+        def record_metric_scope(name: str, metrics: dict[str, float]) -> Iterator[None]:
+            scopes.append(name)
+            try:
+                with timing_metric(name, metrics):
+                    yield
+            finally:
+                scopes.pop()
 
-        self.assertEqual(
-            model_engine.metrics.keys(),
-            {
-                "ctx_cuda_graph_capture_seconds",
-                "post_ctx_cuda_graph_capture_warmup_seconds",
-            },
-        )
-        self.assertTrue(all(value >= 0 for value in model_engine.metrics.values()))
+        def forward(batch: tuple[int, bool], **kwargs: object) -> torch.Tensor:
+            events.append((batch, tuple(scopes)))
+            return torch.empty(1)
+
+        with (
+            patch.object(model_engine, "no_cuda_graph", side_effect=contextlib.nullcontext),
+            patch.object(
+                model_engine,
+                "_create_warmup_request",
+                side_effect=lambda rm, tokens, gen, least_requests=True: (tokens, least_requests),
+            ),
+            patch.object(
+                model_engine,
+                "_release_batch_context",
+                side_effect=lambda batch, rm: contextlib.nullcontext(batch),
+            ),
+            patch.object(model_engine, "_assert_all_tp_ranks_have_warmup_batch"),
+            patch.object(model_engine, "forward", side_effect=forward),
+            patch(
+                "tensorrt_llm._torch.pyexecutor.model_engine.timing_metric",
+                side_effect=record_metric_scope,
+            ),
+            patch(
+                "tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner.timing_metric",
+                side_effect=record_metric_scope,
+            ),
+            patch("tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner.BreakableCUDAGraph"),
+            patch(
+                "tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner.make_weak_ref",
+                side_effect=lambda value: value,
+            ),
+            patch("torch.cuda.Stream"),
+            patch("torch.cuda.current_stream"),
+            patch("torch.cuda.stream", side_effect=lambda stream: contextlib.nullcontext()),
+            patch("torch.cuda.graph_pool_handle", return_value=(1, 2)),
+            patch(
+                "torch.cuda.synchronize", side_effect=lambda: events.append(("sync", tuple(scopes)))
+            ),
+            patch("torch.cuda.empty_cache"),
+            patch("gc.collect"),
+            patch("time.perf_counter", side_effect=iter(range(100))),
+        ):
+            for backend in (PrefillCudaGraphBackend.PIECEWISE, PrefillCudaGraphBackend.BREAKABLE):
+                with self.subTest(backend=backend):
+                    events.clear()
+                    model_engine._metrics.clear()
+                    model_engine.prefill_cuda_graph_backend = backend
+                    model_engine._torch_compile_piecewise_cuda_graph = (
+                        backend == PrefillCudaGraphBackend.PIECEWISE
+                    )
+                    model_engine.breakable_cuda_graph_runner = (
+                        BreakableCUDAGraphRunner(torch.nn.Identity())
+                        if backend == PrefillCudaGraphBackend.BREAKABLE
+                        else None
+                    )
+                    model_engine._capture_prefill_cuda_graphs(resource_manager)
+                    warmup_steps = 3 if backend == PrefillCudaGraphBackend.PIECEWISE else 2
+                    warmup_scope = (
+                        BreakableCUDAGraphRunner.CUDA_GRAPH_WARMUP_METRIC
+                        if backend == PrefillCudaGraphBackend.BREAKABLE
+                        else warmup_metric
+                    )
+                    capture_scope = (
+                        BreakableCUDAGraphRunner.CUDA_GRAPH_CAPTURE_METRIC
+                        if backend == PrefillCudaGraphBackend.BREAKABLE
+                        else capture_metric
+                    )
+                    expected = []
+                    for tokens in (8, 4):
+                        expected.extend([((tokens, True), (warmup_scope,))] * warmup_steps)
+                        expected.extend(
+                            [
+                                ("sync", (warmup_scope,)),
+                                ((tokens, True), (capture_scope,)),
+                                ("sync", (capture_scope,)),
+                            ]
+                        )
+                    for tokens in (8, 4):
+                        expected.extend(
+                            [((tokens, False), (post_metric,)), ("sync", (post_metric,))]
+                        )
+                    self.assertEqual(events, expected)
+                    self.assertEqual(
+                        model_engine.metrics.keys(), {warmup_metric, capture_metric, post_metric}
+                    )
+                    self.assertTrue(all(value >= 0 for value in model_engine.metrics.values()))
+                    self.assertEqual(model_engine.metrics[warmup_metric], 2.0)
+                    self.assertEqual(model_engine.metrics[capture_metric], 2.0)
+                    if backend == PrefillCudaGraphBackend.BREAKABLE:
+                        runner = model_engine.breakable_cuda_graph_runner
+                        self.assertEqual(runner.metrics, {warmup_scope: 1.0, capture_scope: 1.0})
+                        self.assertIsNot(runner.metrics, model_engine.metrics)
+
+                        runner.clear()
+                        model_engine._metrics.clear()
+                        with patch.object(
+                            model_engine, "forward", side_effect=RuntimeError("warmup failed")
+                        ):
+                            with self.assertRaisesRegex(RuntimeError, "warmup failed"):
+                                model_engine._capture_prefill_cuda_graphs(resource_manager)
+                        self.assertEqual(runner.metrics, {warmup_scope: 1.0})
+                        self.assertEqual(model_engine.metrics, {})
+                        self.assertFalse(runner.is_warming_up)
+                        self.assertFalse(runner.is_capturing)
 
     def test_empty_cache_fires_immediately_after_autotuner(self):
         """Change 1 placement: empty_cache must be the call right after

From 744b704305624cc3f5cd91fe3edd65d8c0841a53 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Tue, 22 Sep 2026 16:49:20 -0700
Subject: [PATCH 14/23] imrpve tests

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 ...y_executor_creator_mla_cache_reuse_sync.py | 29 +++++++++++++++++--
 .../test_pytorch_model_engine_warmup.py       | 28 +++++++++++++++++-
 tests/unittest/executor/test_base_worker.py   |  1 +
 3 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
index 0eed05c0cafd..f21d8604995f 100644
--- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
+++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
@@ -19,6 +19,7 @@
 import pytest
 
 from tensorrt_llm._torch.pyexecutor import py_executor_creator
+from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
 from tensorrt_llm._torch.pyexecutor.py_executor_creator import (
     _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS,
     _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS,
@@ -36,6 +37,7 @@
 
 
 def test_executor_creation_stage_metric_names_are_descriptive():
+    """Verify executor creation stages use the expected public metric names."""
     assert py_executor_creator._ExecutorMemoryMonitor.creation_stage_metric_names == {
         ExecutorMemoryType.SAMPLER: "sampler_creation_seconds",
         ExecutorMemoryType.DRAFTER: "speculative_drafter_creation_seconds",
@@ -383,8 +385,10 @@ def _create_py_executor_instance(**kwargs):
     )
 
 
-def test_move_model_engine_metrics_moves_and_clears():
-    py_executor = SimpleNamespace(metrics={})
+def test_move_model_engine_metrics_moves_and_clears() -> None:
+    """Verify main and draft engine metrics are exposed by the executor and cleared at source."""
+    py_executor = object.__new__(PyExecutor)
+    py_executor._metrics = {}
     model_engine = SimpleNamespace(metrics={"total_warmup_seconds": 2.5})
     draft_model_engine = SimpleNamespace(metrics={"total_warmup_seconds": 1.5})
 
@@ -410,7 +414,8 @@ def test_move_model_engine_metrics_moves_and_clears():
     assert draft_model_engine.metrics == {}
 
 
-def test_total_py_executor_creation_metric_finishes_on_error(monkeypatch):
+def test_total_py_executor_creation_metric_finishes_on_success_and_error(monkeypatch) -> None:
+    """Verify creation timing finishes on errors and completed metrics attach on success."""
     captured_metrics = None
 
     @contextmanager
@@ -433,6 +438,24 @@ def _raise_during_creation(**kwargs):
 
     assert captured_metrics == {"total_py_executor_creation_seconds": 1.5}
 
+    py_executor = object.__new__(PyExecutor)
+    py_executor._metrics = {"worker_start_seconds": 0.25}
+
+    def _return_executor(*, creation_metrics: dict[str, float], **kwargs: object) -> PyExecutor:
+        assert "total_py_executor_creation_seconds" not in creation_metrics
+        creation_metrics["model_engine_creation_seconds"] = 1.0
+        return py_executor
+
+    monkeypatch.setattr(py_executor_creator, "_create_py_executor", _return_executor)
+    result = py_executor_creator.create_py_executor(SimpleNamespace())
+
+    assert result is py_executor
+    assert result.metrics == {
+        "worker_start_seconds": 0.25,
+        "model_engine_creation_seconds": 1.0,
+        "total_py_executor_creation_seconds": 1.5,
+    }
+
 
 def test_mla_unsupported_sm_fallback_syncs_cache_reuse(monkeypatch):
     """Verify MLA unsupported SM fallback disables cache reuse in both config and runtime.
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index c438915390d4..c58939ab7543 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -676,7 +676,9 @@ def _record(*msg):
 class TestWarmupCleanup(unittest.TestCase):
     """Lock in warmup-cleanup behavior introduced by PR #14609 (Plan B)."""
 
+    @pytest.mark.cpu_only
     def test_no_kv_cache_warmup_delegates_runner_lifecycle(self):
+        """Verify no-KV-cache warmup delegates warmup and graph capture to the runner."""
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))
         model_engine.moe_load_balancer = None
@@ -699,7 +701,9 @@ def test_no_kv_cache_warmup_delegates_runner_lifecycle(self):
         )
         warmup_sampling.assert_not_called()
 
+    @pytest.mark.cpu_only
     def test_no_kv_cache_warmup_rejects_allocated_kv_cache(self):
+        """Verify a no-KV-cache runner rejects an allocated KV cache before warmup."""
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine._metrics = {}
         model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))
@@ -723,7 +727,8 @@ def test_no_kv_cache_warmup_rejects_allocated_kv_cache(self):
         model_engine._runner.warmup.assert_called_once_with(resource_manager)
 
     @pytest.mark.cpu_only
-    def test_legacy_warmup_skips_without_kv_cache(self) -> None:
+    def test_legacy_warmup_sampling_and_kv_cache_cleanup(self) -> None:
+        """Verify sampling warmup coverage and KV-cache cleanup timing within total warmup."""
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine._warmup_timer = _WarmupTimer(rank=0)
         model_engine._metrics = {}
@@ -790,6 +795,25 @@ def record_metric_scope(name: str, metrics: dict[str, float]) -> Iterator[None]:
         )
         self.assertGreaterEqual(model_engine.metrics["sampling_warmup_seconds"], 0)
 
+        kv_cache_manager = Mock()
+        kv_cache_manager.check_invalid_values_in_kv_cache.return_value = False
+        resource_manager.get_resource_manager.return_value = kv_cache_manager
+        model_engine._metrics.clear()
+        with (
+            patch.object(model_engine, "_warmup_impl") as warmup_impl,
+            patch("time.perf_counter", side_effect=[0.0, 2.0, 5.0, 7.0]),
+        ):
+            model_engine.warmup(resource_manager)
+
+        warmup_impl.assert_called_once_with(resource_manager)
+        kv_cache_manager.check_invalid_values_in_kv_cache.assert_called_once_with(
+            fill_with_zero=True
+        )
+        self.assertEqual(
+            model_engine.metrics,
+            {"kv_cache_cleanup_seconds": 3.0, "total_warmup_seconds": 7.0},
+        )
+
     def test_encoder_decoder_encoder_warmup_delegates_runner_lifecycle(self):
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.model = SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=True))
@@ -807,6 +831,7 @@ def test_encoder_decoder_encoder_warmup_delegates_runner_lifecycle(self):
 
     @pytest.mark.cpu_only
     def test_cuda_graph_metrics_exclude_piecewise_stages(self) -> None:
+        """Verify generation graph timing excludes prefill and includes LoRA warmup cleanup."""
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.model = SimpleNamespace(modules=lambda: [])
         model_engine.llm_args = SimpleNamespace(enable_autotuner=True)
@@ -936,6 +961,7 @@ def record_metric_scope(metric_name: str, _metrics: dict[str, float]) -> Iterato
 
     @pytest.mark.cpu_only
     def test_piecewise_cuda_graph_metrics_are_recorded_separately(self) -> None:
+        """Verify both prefill backends separately time warmup, capture, and post-capture work."""
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE
         model_engine._torch_compile_enabled = True
diff --git a/tests/unittest/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py
index de24dc6fd0ae..dfd8c2f5dde3 100644
--- a/tests/unittest/executor/test_base_worker.py
+++ b/tests/unittest/executor/test_base_worker.py
@@ -26,6 +26,7 @@
 
 @pytest.mark.cpu_only
 def test_get_startup_metrics_promotes_model_engine_stages() -> None:
+    """Verify startup metrics promote saved engine stages while retaining executor metrics."""
     worker = object.__new__(BaseWorker)
     worker._is_pytorch_backend = True
     worker.engine = SimpleNamespace(

From 6cc77a750033595fc2c3ba7b034e36eb61ddba7c Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 23 Sep 2026 12:58:50 -0700
Subject: [PATCH 15/23] minor fixes

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         | 31 +++++++++----------
 .../_torch/pyexecutor/py_executor_creator.py  |  5 ++-
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 6d8e1391ecff..4884ffb01e7d 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3071,19 +3071,21 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager):
                     )
                     if self.breakable_cuda_graph_runner is not None:
                         runner = self.breakable_cuda_graph_runner
-                        runner.capture(
-                            num_tokens, lambda: self.forward(
-                                batch,
-                                new_tensors_device=None,
-                                resource_manager=resource_manager))
-                        self._metrics[
-                            "ctx_cuda_graph_warmup_seconds"] += runner.metrics.get(
-                                BreakableCUDAGraphRunner.
-                                CUDA_GRAPH_WARMUP_METRIC, 0.0)
-                        self._metrics[
-                            "ctx_cuda_graph_capture_seconds"] += runner.metrics.get(
-                                BreakableCUDAGraphRunner.
-                                CUDA_GRAPH_CAPTURE_METRIC, 0.0)
+                        try:
+                            runner.capture(
+                                num_tokens, lambda: self.forward(
+                                    batch,
+                                    new_tensors_device=None,
+                                    resource_manager=resource_manager))
+                        finally:
+                            self._metrics[
+                                "ctx_cuda_graph_warmup_seconds"] += runner.metrics.get(
+                                    BreakableCUDAGraphRunner.
+                                    CUDA_GRAPH_WARMUP_METRIC, 0.0)
+                            self._metrics[
+                                "ctx_cuda_graph_capture_seconds"] += runner.metrics.get(
+                                    BreakableCUDAGraphRunner.
+                                    CUDA_GRAPH_CAPTURE_METRIC, 0.0)
                     else:
                         with timing_metric("ctx_cuda_graph_warmup_seconds",
                                            self._metrics):
@@ -3099,9 +3101,6 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager):
                                          resource_manager=resource_manager)
                             torch.cuda.synchronize()
 
-                    gc.collect()
-                    torch.cuda.empty_cache()
-
         # The logits allocations grow with the number of requests and are not
         # part of the captured model body. Warm up the largest request count so
         # those allocations can be reused during stable inference.
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 3e19b3bac34f..0ed7bb8fae58 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -377,7 +377,9 @@ def _move_model_engine_metrics(
         model_engine: PyTorchModelEngine,
         stage: str,
         draft_model_engine: Optional[PyTorchModelEngine] = None) -> None:
-    """Move model-engine metrics into the executor under a startup stage."""
+    """Move model-engine metrics and optional draft-model-engine metrics into the
+    executor under a startup stage.
+    """
     py_executor.metrics[f"{stage}_model_engine"] = dict(model_engine.metrics)
     model_engine.metrics.clear()
 
@@ -1121,6 +1123,7 @@ def allocation_scope(current_stage: ExecutorMemoryType):
         assert kv_cache_creator is not None
         _move_model_engine_metrics(py_executor, model_engine, "initial",
                                    draft_model_engine)
+        # record initial metrics as py_executor will be deleted later
         initial_model_engine_metrics = dict(py_executor.metrics)
         with allocation_scope(ExecutorMemoryType.MODEL_EXTRA):
             kv_cache_creator.configure_kv_cache_capacity(py_executor)

From a1ebbd4182a244919975e844d533ea616c0abb8e Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 23 Sep 2026 13:12:25 -0700
Subject: [PATCH 16/23] fix test

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../_torch/executor/test_pytorch_model_engine_warmup.py       | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index c58939ab7543..54afe7798c24 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -1085,7 +1085,9 @@ def forward(batch: tuple[int, bool], **kwargs: object) -> torch.Tensor:
                             with self.assertRaisesRegex(RuntimeError, "warmup failed"):
                                 model_engine._capture_prefill_cuda_graphs(resource_manager)
                         self.assertEqual(runner.metrics, {warmup_scope: 1.0})
-                        self.assertEqual(model_engine.metrics, {})
+                        self.assertEqual(
+                            model_engine.metrics, {warmup_metric: 1.0, capture_metric: 0.0}
+                        )
                         self.assertFalse(runner.is_warming_up)
                         self.assertFalse(runner.is_capturing)
 

From 676e3c4108530013413430ca45c6caf59098bd1e Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Thu, 24 Sep 2026 13:15:23 -0700
Subject: [PATCH 17/23] remove difficult-to-maintain test

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../test_pytorch_model_engine_warmup.py       | 135 +-----------------
 1 file changed, 1 insertion(+), 134 deletions(-)

diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 54afe7798c24..6c44ad45c73a 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -14,7 +14,7 @@
 import os
 import sys
 import unittest
-from collections import OrderedDict, defaultdict
+from collections import OrderedDict
 from collections.abc import Iterator
 from dataclasses import dataclass
 from types import ModuleType, SimpleNamespace
@@ -31,7 +31,6 @@
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
 from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM, timing_metric
 from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod
-from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner import BreakableCUDAGraphRunner
 from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller
 from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner
 from tensorrt_llm._torch.pyexecutor.engine.runners.no_kv_cache import NoKVCacheRunner
@@ -959,138 +958,6 @@ def record_metric_scope(metric_name: str, _metrics: dict[str, float]) -> Iterato
             generation.assert_not_called()
             piecewise.assert_not_called()
 
-    @pytest.mark.cpu_only
-    def test_piecewise_cuda_graph_metrics_are_recorded_separately(self) -> None:
-        """Verify both prefill backends separately time warmup, capture, and post-capture work."""
-        model_engine = object.__new__(PyTorchModelEngine)
-        model_engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.PIECEWISE
-        model_engine._torch_compile_enabled = True
-        model_engine._torch_compile_piecewise_cuda_graph = True
-        model_engine._prefill_cuda_graph_num_tokens = [4, 8]
-        model_engine._metrics = defaultdict(float)
-        resource_manager = object()
-        events = []
-        scopes = []
-        warmup_metric = "ctx_cuda_graph_warmup_seconds"
-        capture_metric = "ctx_cuda_graph_capture_seconds"
-        post_metric = "post_ctx_cuda_graph_capture_warmup_seconds"
-
-        @contextlib.contextmanager
-        def record_metric_scope(name: str, metrics: dict[str, float]) -> Iterator[None]:
-            scopes.append(name)
-            try:
-                with timing_metric(name, metrics):
-                    yield
-            finally:
-                scopes.pop()
-
-        def forward(batch: tuple[int, bool], **kwargs: object) -> torch.Tensor:
-            events.append((batch, tuple(scopes)))
-            return torch.empty(1)
-
-        with (
-            patch.object(model_engine, "no_cuda_graph", side_effect=contextlib.nullcontext),
-            patch.object(
-                model_engine,
-                "_create_warmup_request",
-                side_effect=lambda rm, tokens, gen, least_requests=True: (tokens, least_requests),
-            ),
-            patch.object(
-                model_engine,
-                "_release_batch_context",
-                side_effect=lambda batch, rm: contextlib.nullcontext(batch),
-            ),
-            patch.object(model_engine, "_assert_all_tp_ranks_have_warmup_batch"),
-            patch.object(model_engine, "forward", side_effect=forward),
-            patch(
-                "tensorrt_llm._torch.pyexecutor.model_engine.timing_metric",
-                side_effect=record_metric_scope,
-            ),
-            patch(
-                "tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner.timing_metric",
-                side_effect=record_metric_scope,
-            ),
-            patch("tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner.BreakableCUDAGraph"),
-            patch(
-                "tensorrt_llm._torch.pyexecutor.breakable_cuda_graph_runner.make_weak_ref",
-                side_effect=lambda value: value,
-            ),
-            patch("torch.cuda.Stream"),
-            patch("torch.cuda.current_stream"),
-            patch("torch.cuda.stream", side_effect=lambda stream: contextlib.nullcontext()),
-            patch("torch.cuda.graph_pool_handle", return_value=(1, 2)),
-            patch(
-                "torch.cuda.synchronize", side_effect=lambda: events.append(("sync", tuple(scopes)))
-            ),
-            patch("torch.cuda.empty_cache"),
-            patch("gc.collect"),
-            patch("time.perf_counter", side_effect=iter(range(100))),
-        ):
-            for backend in (PrefillCudaGraphBackend.PIECEWISE, PrefillCudaGraphBackend.BREAKABLE):
-                with self.subTest(backend=backend):
-                    events.clear()
-                    model_engine._metrics.clear()
-                    model_engine.prefill_cuda_graph_backend = backend
-                    model_engine._torch_compile_piecewise_cuda_graph = (
-                        backend == PrefillCudaGraphBackend.PIECEWISE
-                    )
-                    model_engine.breakable_cuda_graph_runner = (
-                        BreakableCUDAGraphRunner(torch.nn.Identity())
-                        if backend == PrefillCudaGraphBackend.BREAKABLE
-                        else None
-                    )
-                    model_engine._capture_prefill_cuda_graphs(resource_manager)
-                    warmup_steps = 3 if backend == PrefillCudaGraphBackend.PIECEWISE else 2
-                    warmup_scope = (
-                        BreakableCUDAGraphRunner.CUDA_GRAPH_WARMUP_METRIC
-                        if backend == PrefillCudaGraphBackend.BREAKABLE
-                        else warmup_metric
-                    )
-                    capture_scope = (
-                        BreakableCUDAGraphRunner.CUDA_GRAPH_CAPTURE_METRIC
-                        if backend == PrefillCudaGraphBackend.BREAKABLE
-                        else capture_metric
-                    )
-                    expected = []
-                    for tokens in (8, 4):
-                        expected.extend([((tokens, True), (warmup_scope,))] * warmup_steps)
-                        expected.extend(
-                            [
-                                ("sync", (warmup_scope,)),
-                                ((tokens, True), (capture_scope,)),
-                                ("sync", (capture_scope,)),
-                            ]
-                        )
-                    for tokens in (8, 4):
-                        expected.extend(
-                            [((tokens, False), (post_metric,)), ("sync", (post_metric,))]
-                        )
-                    self.assertEqual(events, expected)
-                    self.assertEqual(
-                        model_engine.metrics.keys(), {warmup_metric, capture_metric, post_metric}
-                    )
-                    self.assertTrue(all(value >= 0 for value in model_engine.metrics.values()))
-                    self.assertEqual(model_engine.metrics[warmup_metric], 2.0)
-                    self.assertEqual(model_engine.metrics[capture_metric], 2.0)
-                    if backend == PrefillCudaGraphBackend.BREAKABLE:
-                        runner = model_engine.breakable_cuda_graph_runner
-                        self.assertEqual(runner.metrics, {warmup_scope: 1.0, capture_scope: 1.0})
-                        self.assertIsNot(runner.metrics, model_engine.metrics)
-
-                        runner.clear()
-                        model_engine._metrics.clear()
-                        with patch.object(
-                            model_engine, "forward", side_effect=RuntimeError("warmup failed")
-                        ):
-                            with self.assertRaisesRegex(RuntimeError, "warmup failed"):
-                                model_engine._capture_prefill_cuda_graphs(resource_manager)
-                        self.assertEqual(runner.metrics, {warmup_scope: 1.0})
-                        self.assertEqual(
-                            model_engine.metrics, {warmup_metric: 1.0, capture_metric: 0.0}
-                        )
-                        self.assertFalse(runner.is_warming_up)
-                        self.assertFalse(runner.is_capturing)
-
     def test_empty_cache_fires_immediately_after_autotuner(self):
         """Change 1 placement: empty_cache must be the call right after
         _run_autotuner_warmup."""

From b751900cdce2f14d8a709004552d8790ee2f7e6e Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Thu, 24 Sep 2026 13:22:16 -0700
Subject: [PATCH 18/23] Format startup timing rebase resolutions

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/model_engine.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 4884ffb01e7d..7c9dbff72935 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1525,7 +1525,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
                 warmup_requests_configs = self._agree_warmup_shapes(
                     self._get_full_general_warmup_requests(resource_manager))
                 # Currently graph has not been captured, disable cuda graph for this warmup.
-                with timing_metric("general_warmup_seconds", self._metrics), self.no_cuda_graph():
+                with timing_metric("general_warmup_seconds",
+                                   self._metrics), self.no_cuda_graph():
                     self._general_warmup(resource_manager,
                                          warmup_requests_configs)
                     # Release C++ MoE workspace buffers so the autotuner can
@@ -1593,7 +1594,8 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             with timing_metric("dg_paged_mqa_warmup_seconds", self._metrics):
                 self._warmup_dg_paged_mqa_logits_metadata()
             log_mem_snapshot("warmup/after_dg_paged_mqa_logits_metadata")
-            with timing_metric("cute_dsl_radix_topk_warmup_seconds", self._metrics):
+            with timing_metric("cute_dsl_radix_topk_warmup_seconds",
+                               self._metrics):
                 self._warmup_cute_dsl_radix_topk()
         log_mem_snapshot("warmup/after_cute_dsl_radix_topk")
         if can_run_general_warmup:
@@ -1602,8 +1604,10 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             with self._warmup_timer.phase("memory_pool_prepop"):
                 warmup_requests_configs = self._get_max_shape_warmup_requests(
                     resource_manager)
-                with timing_metric("memory_pool_prepopulation_seconds", self._metrics):
-                    self._general_warmup(resource_manager, warmup_requests_configs)
+                with timing_metric("memory_pool_prepopulation_seconds",
+                                   self._metrics):
+                    self._general_warmup(resource_manager,
+                                         warmup_requests_configs)
             log_mem_snapshot("warmup/after_memory_pool_prepop")
 
         # Allocate the CUDA graph padding dummies now, while the KV cache is

From 06f1b055394f71cc7cce6e571295e62f6c48a97f Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Mon, 28 Sep 2026 14:26:38 -0700
Subject: [PATCH 19/23] update doc

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 docs/source/llm-api/index.md | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md
index d6fa516bf1e4..b04a4e3b5c2e 100644
--- a/docs/source/llm-api/index.md
+++ b/docs/source/llm-api/index.md
@@ -69,6 +69,9 @@ A typical result has the following structure (values are illustrative):
     "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,
@@ -115,6 +118,9 @@ A typical result has the following structure (values are illustrative):
 }
 ```
 
+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 |
@@ -150,7 +156,7 @@ stages.
 |---------------------|-------|
 | `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` | Warm up general input shapes and release temporary workspaces. |
+| `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. |
@@ -160,7 +166,7 @@ stages.
 | `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` | Pre-populate the memory pool with maximum-shape allocations. |
+| `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. |
 

From 62d9ca982726e1c3281c6a48aaec4ca1ef48863d Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Mon, 28 Sep 2026 14:48:11 -0700
Subject: [PATCH 20/23] merge the work from PR 19402

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_startup.py                      | 13 +++++-
 .../_torch/pyexecutor/model_engine.py         | 40 +++++++++++--------
 .../_torch/pyexecutor/py_executor_creator.py  | 11 ++---
 .../_torch/pyexecutor/warmup_timer.py         | 15 ++++++-
 .../test_pytorch_model_engine_warmup.py       | 18 +++++----
 .../_torch/executor/test_warmup_timer.py      | 13 ++++--
 tests/unittest/test_startup.py                | 24 +++++++----
 7 files changed, 93 insertions(+), 41 deletions(-)

diff --git a/tensorrt_llm/_startup.py b/tensorrt_llm/_startup.py
index af23674158b3..eb8a73d68ea1 100644
--- a/tensorrt_llm/_startup.py
+++ b/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"
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 7c9dbff72935..379f0b9b5bdb 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1466,8 +1466,9 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
         # cuda_graph_config=None flashinfer's sampling kernels would be
         # JIT-built mid-serving.
         self._eager_workspace_reclaimer = None
-        with self._warmup_timer.phase("sampling_module_prewarm"), timing_metric(
-                "sampling_warmup_seconds", self._metrics):
+        with self._warmup_timer.phase("sampling_module_prewarm",
+                                      metrics=self._metrics,
+                                      metric_name="sampling_warmup_seconds"):
             warmup_sampling_module()
             if self.enable_in_graph_sampling:
                 # The fast tier samples inside the captured graph via a
@@ -1514,19 +1515,22 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             self._prewarm_cute_dsl_indexer_q()
         log_mem_snapshot("warmup/after_cute_dsl_indexer_q")
         if not is_enc_dec:
-            with self._warmup_timer.phase("attention_jit"), timing_metric(
-                    "attention_warmup_seconds", self._metrics):
+            with self._warmup_timer.phase(
+                    "attention_jit",
+                    metrics=self._metrics,
+                    metric_name="attention_warmup_seconds"):
                 self._run_attention_warmup(resource_manager,
                                            can_run_general_warmup)
 
         if can_run_general_warmup:
             # Specialize torch.compile graphs across the key input shapes before CUDA graph capture.
-            with self._warmup_timer.phase("general"):
+            with self._warmup_timer.phase("general",
+                                          metrics=self._metrics,
+                                          metric_name="general_warmup_seconds"):
                 warmup_requests_configs = self._agree_warmup_shapes(
                     self._get_full_general_warmup_requests(resource_manager))
                 # Currently graph has not been captured, disable cuda graph for this warmup.
-                with timing_metric("general_warmup_seconds",
-                                   self._metrics), self.no_cuda_graph():
+                with self.no_cuda_graph():
                     self._general_warmup(resource_manager,
                                          warmup_requests_configs)
                     # Release C++ MoE workspace buffers so the autotuner can
@@ -1541,8 +1545,10 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
         # Helix CP is decode-only and runs into issues with the
         # autotuner warmup's context requests.
         if not is_enc_dec and not self.mapping.has_cp_helix():
-            with self._warmup_timer.phase("autotuner"), timing_metric(
-                    "autotuner_warmup_seconds", self._metrics):
+            with self._warmup_timer.phase(
+                    "autotuner",
+                    metrics=self._metrics,
+                    metric_name="autotuner_warmup_seconds"):
                 self._run_autotuner_warmup(resource_manager)
             log_mem_snapshot("warmup/after_autotuner")
             # Pre-JIT Mamba SSD multi-seq + HAS_INITSTATES=True Triton kernels
@@ -1550,8 +1556,10 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
             # since MambaHybridCacheManager skips _general_warmup and the
             # default autotuner shape is single-seq / no-initstates. Safe
             # no-op for non-Mamba models.
-            with self._warmup_timer.phase("mamba_hybrid"), timing_metric(
-                    "mamba_hybrid_warmup_seconds", self._metrics):
+            with self._warmup_timer.phase(
+                    "mamba_hybrid",
+                    metrics=self._metrics,
+                    metric_name="mamba_hybrid_warmup_seconds"):
                 self._run_mamba_hybrid_warmup(resource_manager)
             log_mem_snapshot("warmup/after_mamba_hybrid")
             # Release the autotuner's exploration-mode intermediates. The
@@ -1601,13 +1609,13 @@ def _warmup_scheduled(self, resource_manager: ResourceManager,
         if can_run_general_warmup:
             # Pre-populate the memory pool with max-shape allocations to reduce
             # fragmentation at runtime.
-            with self._warmup_timer.phase("memory_pool_prepop"):
+            with self._warmup_timer.phase(
+                    "memory_pool_prepop",
+                    metrics=self._metrics,
+                    metric_name="memory_pool_prepopulation_seconds"):
                 warmup_requests_configs = self._get_max_shape_warmup_requests(
                     resource_manager)
-                with timing_metric("memory_pool_prepopulation_seconds",
-                                   self._metrics):
-                    self._general_warmup(resource_manager,
-                                         warmup_requests_configs)
+                self._general_warmup(resource_manager, warmup_requests_configs)
             log_mem_snapshot("warmup/after_memory_pool_prepop")
 
         # Allocate the CUDA graph padding dummies now, while the KV cache is
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 0ed7bb8fae58..80b1ca1b6e68 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -676,9 +676,9 @@ def allocation_scope(current_stage: ExecutorMemoryType):
         }.get(current_stage, current_stage.value)
         stage = current_stage.value
         metric_name = mem_monitor.creation_stage_metric_names[current_stage]
-        with _startup_timer.phase(timing_name), timing_metric(
-                metric_name,
-                creation_metrics), mem_monitor.observe_creation_stage(
+        with _startup_timer.phase(
+                timing_name, metrics=creation_metrics,
+                metric_name=metric_name), mem_monitor.observe_creation_stage(
                     current_stage):
             if not enable_sleep or stage.startswith("_no_capture"):
                 yield
@@ -1202,8 +1202,9 @@ def allocation_scope(current_stage: ExecutorMemoryType):
     if mapping.rank == 0:
         logger.info(f"LLM Args:\n{llm_args}")
 
-    with _startup_timer.phase("executor_start_worker"), timing_metric(
-            "worker_start_seconds", creation_metrics):
+    with _startup_timer.phase("executor_start_worker",
+                              metrics=creation_metrics,
+                              metric_name="worker_start_seconds"):
         py_executor.start_worker()
 
     py_executor.metrics.update(initial_model_engine_metrics)
diff --git a/tensorrt_llm/_torch/pyexecutor/warmup_timer.py b/tensorrt_llm/_torch/pyexecutor/warmup_timer.py
index 96101b7f7328..2fedebd6a557 100644
--- a/tensorrt_llm/_torch/pyexecutor/warmup_timer.py
+++ b/tensorrt_llm/_torch/pyexecutor/warmup_timer.py
@@ -43,11 +43,21 @@ def prefix(self) -> str:
         )
 
     @contextmanager
-    def phase(self, name: str, *, record: bool = True, log_start: bool = True) -> Iterator[None]:
+    def phase(
+        self,
+        name: str,
+        *,
+        record: bool = True,
+        log_start: bool = True,
+        metrics: dict[str, float] | None = None,
+        metric_name: str | None = None,
+    ) -> Iterator[None]:
         """Log start/end of one warmup phase and record its wall-clock time.
 
         ``name`` should match the ``log_mem_snapshot`` tag of the same phase
         (``warmup/after_<name>``) so timing and memory logs line up.
+        When ``metrics`` is provided, accumulate the same elapsed value under
+        ``metric_name`` (or ``name``), independently of summary recording.
         """
         if log_start:
             logger.info(f"{self.prefix} {name}: start")
@@ -60,6 +70,9 @@ def phase(self, name: str, *, record: bool = True, log_start: bool = True) -> It
             elapsed = time.perf_counter() - start
             if record:
                 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
             if completed:
                 logger.info(f"{self.prefix} {name}: done in {elapsed:.1f}s")
             else:
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 6c44ad45c73a..24258e2304a2 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -29,7 +29,7 @@
 from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8GemmRunner
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper
-from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM, timing_metric
+from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM
 from tensorrt_llm._torch.modules.linear import MXFP8LinearMethod
 from tensorrt_llm._torch.pyexecutor.engine.model_call import ModelCaller
 from tensorrt_llm._torch.pyexecutor.engine.runners.encoder_decoder import EncoderDecoderRunner
@@ -743,17 +743,21 @@ def test_legacy_warmup_sampling_and_kv_cache_cleanup(self) -> None:
         resource_manager = Mock()
         resource_manager.get_resource_manager.return_value = None
         events = []
+        phase = model_engine._warmup_timer.phase
 
         @contextlib.contextmanager
-        def record_metric_scope(name: str, metrics: dict[str, float]) -> Iterator[None]:
-            events.append(("enter", name))
-            with timing_metric(name, metrics):
+        def record_metric_scope(
+            name: str, *, metrics: dict[str, float], metric_name: str
+        ) -> Iterator[None]:
+            events.append(("enter", metric_name))
+            with phase(name, metrics=metrics, metric_name=metric_name):
                 yield
-            events.append(("exit", name))
+            events.append(("exit", metric_name))
 
         with (
-            patch(
-                "tensorrt_llm._torch.pyexecutor.model_engine.timing_metric",
+            patch.object(
+                model_engine._warmup_timer,
+                "phase",
                 side_effect=record_metric_scope,
             ),
             patch(
diff --git a/tests/unittest/_torch/executor/test_warmup_timer.py b/tests/unittest/_torch/executor/test_warmup_timer.py
index 10dae1fcb5ea..1dd4059b9bd9 100644
--- a/tests/unittest/_torch/executor/test_warmup_timer.py
+++ b/tests/unittest/_torch/executor/test_warmup_timer.py
@@ -30,10 +30,13 @@ def test_warmup_phase_repeated_names(warmup_timing):
     """Repeated names accumulate and the summary reports identity and percentages."""
     engine, clock, log = warmup_timing
     clock.side_effect = [0, 2, 3, 6]
+    metrics = {}
     for _ in range(2):
-        with engine.phase("attention"):
+        with engine.phase("attention", metrics=metrics, metric_name="attention_warmup_seconds"):
             pass
     assert engine.timings == {"attention": 5}
+    assert metrics == {"attention_warmup_seconds": 5}
+    assert clock.call_count == 4
     engine.summary(10)
     log.info.assert_any_call(
         "[warmup][pid=42][rank=3][purpose=memory_profiling][pass=1] "
@@ -47,12 +50,14 @@ def test_warmup_phase_failure_records_elapsed(warmup_timing, error_type):
     """Preserve elapsed time and the original exception on interrupted warmup."""
     engine, clock, log = warmup_timing
     clock.side_effect = [2, 7]
+    metrics = {}
     error = error_type("warmup failed")
     with pytest.raises(error_type) as caught:
-        with engine.phase("attention"):
+        with engine.phase("attention", metrics=metrics):
             raise error
     assert caught.value is error
     assert engine.timings == {"attention": 5}
+    assert metrics == {"attention": 5}
     log.warning.assert_called_once_with(
         "[warmup][pid=42][rank=3][purpose=memory_profiling][pass=1] attention: failed after 5.0s"
     )
@@ -120,11 +125,13 @@ def test_warmup_shapes_are_not_double_counted(warmup_timing, log_start):
     """Shape timing stays visible without inflating the enclosing phase total."""
     _, clock, log = warmup_timing
     clock.side_effect = [0, 1, 2, 5, 7, 10]
+    metrics = {}
     with warmup_timer._WarmupTimer(rank=3) as timer:
         with timer.phase("general"):
-            with timer.phase("general shape", record=False, log_start=log_start):
+            with timer.phase("general shape", record=False, log_start=log_start, metrics=metrics):
                 pass
     assert timer.timings == {"general": 6}
+    assert metrics == {"general shape": 3}
     messages = [call.args[0] for call in log.info.call_args_list]
     assert any("general shape: done in 3.0s" in msg for msg in messages)
     assert any("general shape: start" in msg for msg in messages) == log_start
diff --git a/tests/unittest/test_startup.py b/tests/unittest/test_startup.py
index e9921144944d..0f00b94977b9 100644
--- a/tests/unittest/test_startup.py
+++ b/tests/unittest/test_startup.py
@@ -24,15 +24,18 @@ def timing(monkeypatch):
 
 
 def test_nested_phases_and_unattributed_time(timing):
-    """Count nested time once and retain time outside measured phases."""
+    """Export nested phases without double counting them in the summary."""
     clock, log = timing
     clock.side_effect = [0, 2, 3, 4, 6, 9, 12]
+    metrics = {}
     with _startup._StartupTimer("executor") as timer:
         timer.mark_initialization("configuration")
-        with timer.phase("model"):
-            with timer.phase("weights"):
+        with timer.phase("model", metrics=metrics):
+            with timer.phase("weights", metrics=metrics):
                 pass
     assert timer.timings == {"configuration": 2, "model": 6}
+    assert metrics == {"model": 6, "weights": 2}
+    assert clock.call_count == 7
     assert timer.depth == 0
     log.info.assert_any_call("[startup][pid=42] executor/weights: done in 2.000s")
     log.info.assert_any_call(
@@ -42,15 +45,18 @@ def test_nested_phases_and_unattributed_time(timing):
 
 
 def test_repeated_phases_accumulate(timing):
-    """Keep both occurrences of a repeated phase in the summary."""
+    """Accumulate repeated phases in both the summary and exported metrics."""
     clock, _ = timing
     clock.side_effect = [0, 1, 3, 4, 7, 8]
+    metrics = {}
     with _startup._StartupTimer("executor") as timer:
-        with timer.phase("allocation"):
+        with timer.phase("allocation", metrics=metrics, metric_name="allocation_seconds"):
             pass
-        with timer.phase("allocation"):
+        with timer.phase("allocation", metrics=metrics, metric_name="allocation_seconds"):
             pass
     assert timer.timings == {"allocation": 5}
+    assert metrics == {"allocation_seconds": 5}
+    assert clock.call_count == 6
 
 
 @pytest.mark.parametrize("error_type", [ValueError, KeyboardInterrupt])
@@ -60,14 +66,16 @@ def test_failed_nested_phase_preserves_time_and_exception(timing, error_type):
     clock.side_effect = [0, 1, 2, 4, 6, 8]
     error = error_type("initialization failed")
     timer = _startup._StartupTimer("executor")
+    metrics = {}
     with pytest.raises(error_type) as caught:
         with timer:
-            with timer.phase("model"):
-                with timer.phase("weights"):
+            with timer.phase("model", metrics=metrics):
+                with timer.phase("weights", metrics=metrics):
                     raise error
     assert caught.value is error
     assert timer.depth == 0
     assert timer.timings == {"model": 5}
+    assert metrics == {"model": 5, "weights": 2}
     messages = [c.args[0] for c in log.info.call_args_list]
     assert "[startup][pid=42] executor/weights: failed in 2.000s" in messages
     assert "[startup][pid=42] executor/model: failed in 5.000s" in messages

From 651d6a215dbd89ab8f0946e8638cd9b21e6309a0 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Tue, 29 Sep 2026 12:08:38 -0700
Subject: [PATCH 21/23] fix test

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../unittest/_torch/executor/test_pytorch_model_engine_warmup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 24258e2304a2..623764ca93be 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -344,6 +344,7 @@ def test_compiled_mxfp8_warmup_backend_selection(
     method = MXFP8LinearMethod()
     engine = SimpleNamespace(
         llm_args=SimpleNamespace(enable_autotuner=True),
+        _warmup_timer=_WarmupTimer(rank=0),
         _torch_compile_enabled=compile_enabled,
         _torch_compile_prefill_only=prefill_only,
         _torch_compile_backend=None,

From b311bd9e80856d8fe3410535d78ca355bfb52448 Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Tue, 29 Sep 2026 14:41:59 -0700
Subject: [PATCH 22/23] [TRTLLM-15447][fix] Resolve warmup timer merge
 duplication

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../unittest/_torch/executor/test_pytorch_model_engine_warmup.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 623764ca93be..24258e2304a2 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -344,7 +344,6 @@ def test_compiled_mxfp8_warmup_backend_selection(
     method = MXFP8LinearMethod()
     engine = SimpleNamespace(
         llm_args=SimpleNamespace(enable_autotuner=True),
-        _warmup_timer=_WarmupTimer(rank=0),
         _torch_compile_enabled=compile_enabled,
         _torch_compile_prefill_only=prefill_only,
         _torch_compile_backend=None,

From 7ed62570bb547215a45efcb9468b42193ddbb21b Mon Sep 17 00:00:00 2001
From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Date: Wed, 30 Sep 2026 11:02:54 -0700
Subject: [PATCH 23/23] [TRTLLM-15447][fix] Update warmup fixture for LoRA
 builder

Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
---
 .../_torch/executor/test_pytorch_model_engine_warmup.py         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
index 24258e2304a2..50c876092f12 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py
@@ -838,7 +838,7 @@ def test_cuda_graph_metrics_exclude_piecewise_stages(self) -> None:
         model_engine = object.__new__(PyTorchModelEngine)
         model_engine.model = SimpleNamespace(modules=lambda: [])
         model_engine.llm_args = SimpleNamespace(enable_autotuner=True)
-        model_engine.cuda_graph_lora_manager = object()
+        model_engine._lora = SimpleNamespace(cuda_graph_manager=object())
         model_engine.cuda_graph_runner = SimpleNamespace(
             enabled=True,
             is_warmup_only=False,