diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 2d0a9f0cf64f..b04a4e3b5c2e 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -50,44 +50,143 @@ llm = LLM(model=) ## Startup Metrics -For the PyTorch backend, the beta `LLM.startup_metrics` property reports weight-loading timings from -worker rank 0. Values are wall-clock seconds. The property returns an empty dictionary when the -backend does not provide startup metrics. +For the PyTorch backend, the beta `LLM.startup_metrics` property reports executor construction, +model-engine warmup, and weight-loading timings from worker rank 0. Values are wall-clock seconds. +The property returns an empty dictionary when the backend does not provide startup metrics. ```python llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0") print(llm.startup_metrics) ``` -A typical result has the following structure: +A typical result has the following structure (values are illustrative): ```json { + "initial_model_engine": { + "sampling_warmup_seconds": 0.342, + "attention_warmup_seconds": 2.911, + "general_warmup_seconds": 0.387, + "autotuner_warmup_seconds": 0.041, + "mamba_hybrid_warmup_seconds": 0.000, + "ctx_cuda_graph_warmup_seconds": 0.480, + "ctx_cuda_graph_capture_seconds": 0.320, + "post_ctx_cuda_graph_capture_warmup_seconds": 0.110, + "gen_cuda_graph_warmup_seconds": 0.566, + "gen_cuda_graph_capture_seconds": 0.950, + "dg_paged_mqa_warmup_seconds": 0.000, + "cute_dsl_radix_topk_warmup_seconds": 0.000, + "memory_pool_prepopulation_seconds": 0.094, + "kv_cache_cleanup_seconds": 0.194, + "total_warmup_seconds": 5.485 + }, + "final_model_engine": { + "sampling_warmup_seconds": 0.291, + "attention_warmup_seconds": 0.057, + "general_warmup_seconds": 0.492, + "autotuner_warmup_seconds": 0.046, + "mamba_hybrid_warmup_seconds": 0.000, + "gen_cuda_graph_warmup_seconds": 0.599, + "gen_cuda_graph_capture_seconds": 0.991, + "dg_paged_mqa_warmup_seconds": 0.000, + "cute_dsl_radix_topk_warmup_seconds": 0.000, + "memory_pool_prepopulation_seconds": 0.097, + "kv_cache_cleanup_seconds": 1.657, + "total_warmup_seconds": 4.230 + }, + "py_executor": { + "config_and_checkpoint_loader_initialization_seconds": 0.004, + "model_engine_creation_seconds": 1.554, + "sampler_creation_seconds": 1.195, + "initial_kv_cache_creation_seconds": 0.025, + "speculative_decoding_resource_manager_creation_seconds": 0.000, + "speculative_drafter_creation_seconds": 0.000, + "initial_py_executor_creation_seconds_for_kv_cache_estimation": 5.766, + "kv_cache_capacity_configuration_seconds": 0.075, + "final_kv_cache_creation_seconds": 0.057, + "final_py_executor_creation_seconds": 4.459, + "worker_start_seconds": 0.000, + "total_py_executor_creation_seconds": 13.373 + }, "model_loader": { - "total_model_loading_seconds": 1.971, - "checkpoint_preparation_seconds": 1.177, - "weight_population_seconds": 0.598, + "checkpoint_preparation_seconds": 0.719, + "weight_population_seconds": 0.272, "checkpoint_finalization_seconds": 0.031, - "post_load_processing_seconds": 0.005 + "post_load_processing_seconds": 0.004, + "total_model_loading_seconds": 1.545 } } ``` -The `model_loader` object contains timings for the main LLM weights. If a draft -model is used, the additional fields `draft_checkpoint_preparation_seconds`, -`draft_weight_population_seconds`, and `draft_checkpoint_finalization_seconds` -appear. +Some of the timing metrics may not appear or may show as a very small value if the corresponding codepath is not executed, e.g. +if the prefill CUDA graphs are not enabled, its related metrics "...ctx_cuda_graph..." will not appear. + +The `py_executor` property contains the timed scopes in `create_py_executor()`, which creates the [PyExecutor](../developer-guide/overview.md). + +| PyExecutor metric | Scope | +|-----------------|-------| +| `config_and_checkpoint_loader_initialization_seconds` | Initialize the model config and checkpoint loader. | +| `model_engine_creation_seconds` | Construct the main model engine, which includes loading model weights. | +| `draft_model_engine_creation_seconds` | Construct the separate draft model engine in the deprecated two-model MTP setting. | +| `guided_decoder_creation_seconds` | Construct guided-decoding resources. | +| `sampler_creation_seconds` | Construct the sampler. | +| `initial_kv_cache_creation_seconds` | Construct the temporary KV cache used to estimate final capacity. | +| `speculative_decoding_resource_manager_creation_seconds` | Construct the speculative-decoding resource manager. | +| `speculative_drafter_creation_seconds` | Construct the speculative drafter. | +| `initial_py_executor_creation_seconds_for_kv_cache_estimation` | Construct the temporary PyExecutor for KV cache capacity estimation, which includes the initial model engine warmup. | +| `kv_cache_capacity_configuration_seconds` | Determine final KV cache capacity from the temporary executor and available memory. | +| `final_kv_cache_creation_seconds` | Construct the final KV cache retained for serving. | +| `final_py_executor_creation_seconds` | Using the final KV cache, construct the final PyExecutor retained for serving, which includes the final model engine warmup. | +| `worker_start_seconds` | Start the final PyExecutor worker. | +| `total_py_executor_creation_seconds` | Total time for the `create_py_executor()` call, including all applicable scopes above. | + +The two model-engine properties contain times for various `PyTorchModelEngine` warmup stages. +`initial_model_engine` measures the timings for the initial model engine creation as part of the +`py_executor.initial_py_executor_creation_seconds_for_kv_cache_estimation` timing. +Likewise, `final_model_engine` measures the timings for the final model engine creation as part of +the `py_executor.final_py_executor_creation_seconds` timing. +The deprecated two-model speculative decoding setting may produce the additional properties +`initial_draft_model_engine` and `final_draft_model_engine`. + +Each model-engine property contains the warmup scopes below. Which scopes appear depends on the model +and configuration; for example, encoder-decoder and context-parallel configurations can skip some +stages. + +| Model-engine metric | Scope | +|---------------------|-------| +| `sampling_warmup_seconds` | Warm up sampling kernels and run `torch.compile()` on sampling operators. | +| `attention_warmup_seconds` | Warm up the attention backend and kernels. | +| `general_warmup_seconds` | Prepare and warm up general input shapes and release temporary workspaces. | +| `autotuner_warmup_seconds` | Run kernel autotuning warmup. | +| `mamba_hybrid_warmup_seconds` | Warm up Mamba hybrid kernels, when applicable. | +| `gen_cuda_graph_warmup_seconds` | Run the warmup-only pass for generation and mixed encoder-decoder CUDA graphs, including LoRA autotuning and pipeline-parallel cache synchronization and cleanup. | +| `gen_cuda_graph_capture_seconds` | Capture generation and mixed encoder-decoder CUDA graphs for serving. | +| `ctx_cuda_graph_warmup_seconds` | Warm up prefill forwards before CUDA graph capture. | +| `ctx_cuda_graph_capture_seconds` | Capture prefill CUDA graphs. | +| `post_ctx_cuda_graph_capture_warmup_seconds` | Warm up prefill logits buffers with many requests after CUDA graph capture. | +| `dg_paged_mqa_warmup_seconds` | Warm up DeepGEMM paged-MQA metadata, when applicable. | +| `cute_dsl_radix_topk_warmup_seconds` | Warm up the CuTe DSL radix top-k kernel, when applicable. | +| `memory_pool_prepopulation_seconds` | Prepare maximum-shape requests and pre-populate the memory pool. | +| `kv_cache_cleanup_seconds` | Check and clear invalid KV cache values produced during warmup. | +| `total_warmup_seconds` | Complete model-engine warmup, including KV cache cleanup. | + +The `model_loader` property contains timings for loading the main LLM weights, including +`checkpoint_finalization_seconds` for checkpoint-session synchronization and cleanup after weight population. +If a draft model is used, additional fields `draft_checkpoint_preparation_seconds`, +`draft_weight_population_seconds`, and `draft_checkpoint_finalization_seconds` will appear. +A `draft_model_loader` property can also appear in the deprecated two-model MTP setting. +The `py_executor.model_engine_creation_seconds` timing includes the total `model_loader` timing. | Metric | Description | |--------|-------------| -| `total_model_loading_seconds` | Overall model construction and loading interval measured after checkpoint configuration validation. It includes the named phases below. | | `checkpoint_preparation_seconds` | Time spent warming up, parsing and preparing checkpoint tensors for the model. Some checkpoint formats can populate model storage directly during this phase. | | `weight_population_seconds` | Time spent copying prepared checkpoint tensors into model parameters on GPUs. This metric can be absent for formats that populate weights directly during the above checkpoint preparation phase. | -| `checkpoint_finalization_seconds` | Time spent finalizing the checkpoint session after weight population. This includes loader-specific synchronization and cleanup; rank-striped read-ahead includes waiting for peer ranks and stopping background readers. | +| `checkpoint_finalization_seconds` | Time spent finalizing the checkpoint session after weight population. This includes loader-specific synchronization and cleanup. | | `draft_checkpoint_preparation_seconds` | Checkpoint preparation time for draft weights loaded as part of the model loader. | | `draft_weight_population_seconds` | Weight population time for draft weights loaded as part of the model loader. | | `draft_checkpoint_finalization_seconds` | Checkpoint finalization time for draft weights loaded as part of the model loader. | | `post_load_processing_seconds` | Time spent in format-specific hooks and model finalization, including post-load weight transformation, quantization and memory cleanup. | +| `total_model_loading_seconds` | Overall model construction and loading interval measured after checkpoint configuration validation. It includes the named phases below. | `trtllm-serve` exposes the same rank-0 payload in the `startup_metrics` field of the `GET /server_info` response: @@ -100,7 +199,7 @@ curl http://localhost:8000/server_info { "startup_metrics": { "model_loader": { - "total_model_loading_seconds": 1.971, + "total_model_loading_seconds": 1.545, ... } } 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/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 da57e8c27eda..379f0b9b5bdb 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 @@ -55,7 +57,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 +380,7 @@ def __init__( ): _configure_deep_gemm_pdl() + self._metrics: dict[str, float] = defaultdict(float) self.forward_pass_callable = None self._cleanup_done = False self._model_caller: Optional[ModelCaller] = None @@ -1180,6 +1183,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]]: @@ -1295,17 +1303,19 @@ def wrapper(self, resource_manager: Optional[ResourceManager] = None, *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("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 @@ -1422,6 +1432,12 @@ 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.""" + 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. @@ -1450,7 +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"): + 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 @@ -1497,13 +1515,18 @@ 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", + 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. @@ -1522,7 +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"): + 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 @@ -1530,7 +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"): + 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 @@ -1551,8 +1580,7 @@ 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) + self._run_cuda_graph_warmup(resource_manager) finally: self.cuda_graph_runner.is_warmup_only = False self.cuda_graph_runner.padding_dummy_requests = {} @@ -1571,14 +1599,20 @@ 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 # 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) self._general_warmup(resource_manager, warmup_requests_configs) @@ -2596,34 +2630,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, @@ -3004,7 +3048,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.""" + """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 + 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 @@ -3033,47 +3082,67 @@ def _capture_prefill_cuda_graphs(self, resource_manager: ResourceManager): 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)) + runner = self.breakable_cuda_graph_runner + 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: - # Run a few times to ensure torch.compile capture. - for _ in range(4): + 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() # 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() + 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, 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() ### Helper methods promoted from the original warmup method ### 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..80b1ca1b6e68 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,23 @@ 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 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() + + 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 +417,38 @@ def _create_py_executor_impl( Returns: A fully initialized PyExecutor instance. """ + creation_metrics: dict[str, float] = {} + with timing_metric("total_py_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 +674,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) + stage = current_stage.value + metric_name = mem_monitor.creation_stage_metric_names[current_stage] with _startup_timer.phase( - timing_name), mem_monitor.observe_creation_stage(current_stage): - stage = current_stage.value + 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 else: @@ -1047,6 +1121,10 @@ 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) + # 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) @@ -1124,9 +1202,15 @@ 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", + metrics=creation_metrics, + metric_name="worker_start_seconds"): 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/_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_``) 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/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index c4dfa7552244..af70bc4ed487 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["py_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..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 @@ -13,23 +13,50 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import contextmanager from types import SimpleNamespace 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, ) 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(): + """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", + 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 +106,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 +174,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 +385,78 @@ def _create_py_executor_instance(**kwargs): ) +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}) + + 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( + 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_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 + 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_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 edba96b46000..50c876092f12 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", @@ -654,10 +675,13 @@ 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 + model_engine._metrics = {} model_engine.is_warmup = False model_engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER model_engine._fallback_to_engine = False @@ -676,8 +700,11 @@ 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)) model_engine.moe_load_balancer = None model_engine.is_warmup = False @@ -699,30 +726,96 @@ 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_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 = {} model_engine.moe_load_balancer = None model_engine.is_warmup = False model_engine.enable_in_graph_sampling = False 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 = [] + phase = model_engine._warmup_timer.phase + + @contextlib.contextmanager + 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", metric_name)) with ( + patch.object( + model_engine._warmup_timer, + "phase", + side_effect=record_metric_scope, + ), patch( - "tensorrt_llm._torch.pyexecutor.model_engine.warmup_sampling_module" + "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) + + 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) @@ -739,6 +832,136 @@ 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: + """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) + model_engine._lora = SimpleNamespace(cuda_graph_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_empty_cache_fires_immediately_after_autotuner(self): """Change 1 placement: empty_cache must be the call right after _run_autotuner_warmup.""" 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/executor/test_base_worker.py b/tests/unittest/executor/test_base_worker.py index bc9a814e76c7..dfd8c2f5dde3 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,64 @@ model_path = llm_models_root() / default_model_name +@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( + 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 + }, + "py_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(): 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