Add span metrics for model_forward, scheduler and sampler time

tests/tracing/test_tracing.py

@@ -114,3 +114,5 @@ def test_traces(trace_service):
         SpanAttributes.LLM_LATENCY_TIME_TO_FIRST_TOKEN) == ttft
     e2e_time = metrics.finished_time - metrics.arrival_time
     assert attributes.get(SpanAttributes.LLM_LATENCY_E2E) == e2e_time
+    assert attributes.get(
+        SpanAttributes.LLM_LATENCY_TIME_IN_SCHEDULER) == metrics.scheduler_time_ms

vllm/config.py

@@ -1576,6 +1576,11 @@ class ObservabilityConfig:
     """Configuration for observability."""
     otlp_traces_endpoint: Optional[str] = None
 
+    # If set, collects the model forward time for the request. This introduces a possibly
+    # blocking operation to accurately collect the GPU time. It can have a performance 
+    # impact on the request latency.
+    collect_model_forward_time: bool = False
+
     def __post_init__(self):
         if not is_otel_installed() and self.otlp_traces_endpoint is not None:
             raise ValueError("OpenTelemetry packages must be installed before "

vllm/core/scheduler.py

@@ -975,6 +975,7 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
         # Schedule sequence groups.
         # This function call changes the internal states of the scheduler
         # such as self.running, self.swapped, and self.waiting.
+        scheduler_start_time = time.time()
         scheduler_outputs = self._schedule()
         now = time.time()
 
@@ -1047,7 +1048,16 @@ def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
         for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups:
             self.block_manager.mark_blocks_as_computed(
                 scheduled_seq_group.seq_group)
-
+
+        scheduler_time = time.time() - scheduler_start_time
+        # Add this to scheduler time to all the sequences that are currently running.
+        # This will help estimate if the scheduler is a significant component in the e2e latency.
+        for seq_group in self.running:
+            if seq_group.metrics.scheduler_time is not None:
+                seq_group.metrics.scheduler_time += scheduler_time
+            else:
+                seq_group.metrics.scheduler_time = scheduler_time
+
         return seq_group_metadata_list, scheduler_outputs
 
     def fork_seq(self, parent_seq: Sequence, child_seq: Sequence) -> None:

vllm/engine/arg_utils.py

@@ -117,6 +117,7 @@ class EngineArgs:
     disable_logprobs_during_spec_decoding: Optional[bool] = None
 
     otlp_traces_endpoint: Optional[str] = None
+    collect_model_forward_time: bool = False
 
     def __post_init__(self):
         if self.tokenizer is None:
@@ -660,6 +661,14 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             type=str,
             default=None,
             help='Target URL to which OpenTelemetry traces will be sent.')
+        parser.add_argument(
+            '--collect-model-forward-time',
+             type=bool,
+            default=EngineArgs.collect_model_forward_time,
+            help="If set to True and otlp-traces-endpoint is set, "
+            "collects model forward time in traces. This involves "
+            "use of a blocking operation and hence might have a "
+            "performance impact.")
 
         return parser
 
@@ -845,7 +854,8 @@ def create_engine_config(self, ) -> EngineConfig:
             guided_decoding_backend=self.guided_decoding_backend)
 
         observability_config = ObservabilityConfig(
-            otlp_traces_endpoint=self.otlp_traces_endpoint)
+            otlp_traces_endpoint=self.otlp_traces_endpoint,
+            collect_model_forward_time=self.collect_model_forward_time)
 
         if (model_config.get_sliding_window() is not None
                 and scheduler_config.chunked_prefill_enabled

vllm/engine/llm_engine.py

@@ -259,6 +259,7 @@ def __init__(
             speculative_config=speculative_config,
             load_config=load_config,
             prompt_adapter_config=prompt_adapter_config,
+            observability_config=self.observability_config,
         )
 
         if not self.model_config.embedding_mode:
@@ -814,6 +815,16 @@ def _process_model_outputs(
             seq_group = scheduled_seq_group.seq_group
             seq_group.update_num_computed_tokens(
                 scheduled_seq_group.token_chunk_size)
+            if output is not None and len(output) > 0 and isinstance(output[0], SamplerOutput):
+                for o in output:
+                    if seq_group.metrics.model_forward_time is not None:
+                        seq_group.metrics.model_forward_time += o.model_forward_time
+                    else:
+                        seq_group.metrics.model_forward_time = o.model_forward_time
+                    if seq_group.metrics.model_execute_time is not None:
+                        seq_group.metrics.model_execute_time += o.model_execute_time
+                    else:
+                        seq_group.metrics.model_execute_time = o.model_execute_time
             if self.model_config.embedding_mode:
                 self._process_sequence_group_outputs(seq_group, outputs)
                 continue
@@ -1205,3 +1216,9 @@ def create_trace_span(self, seq_group: SequenceGroup) -> None:
             seq_span.set_attribute(
                 SpanAttributes.LLM_LATENCY_TIME_TO_FIRST_TOKEN, ttft)
             seq_span.set_attribute(SpanAttributes.LLM_LATENCY_E2E, e2e_time)
+            if metrics.scheduler_time is not None:
+                seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_SCHEDULER, metrics.scheduler_time)
+            if metrics.model_forward_time is not None:
+                seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_FORWARD, metrics.model_forward_time / 1000.0)
+            if metrics.model_execute_time is not None:
+                seq_span.set_attribute(SpanAttributes.LLM_LATENCY_TIME_IN_MODEL_EXECUTE, metrics.model_execute_time)

vllm/executor/executor_base.py

@@ -2,8 +2,8 @@
 from typing import List, Optional, Set, Tuple
 
 from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig,
-                         ModelConfig, MultiModalConfig, ParallelConfig,
-                         PromptAdapterConfig, SchedulerConfig,
+                         ModelConfig, MultiModalConfig, ObservabilityConfig,
+                         ParallelConfig, PromptAdapterConfig, SchedulerConfig,
                          SpeculativeConfig)
 from vllm.lora.request import LoRARequest
 from vllm.prompt_adapter.request import PromptAdapterRequest
@@ -32,6 +32,7 @@ def __init__(
         multimodal_config: Optional[MultiModalConfig],
         speculative_config: Optional[SpeculativeConfig],
         prompt_adapter_config: Optional[PromptAdapterConfig],
+        observability_config: Optional[ObservabilityConfig],
     ) -> None:
         self.model_config = model_config
         self.cache_config = cache_config
@@ -43,7 +44,7 @@ def __init__(
         self.multimodal_config = multimodal_config
         self.speculative_config = speculative_config
         self.prompt_adapter_config = prompt_adapter_config
-
+        self.observability_config = observability_config
         self._init_executor()
 
     @abstractmethod

vllm/executor/gpu_executor.py

@@ -60,6 +60,7 @@ def _get_worker_kwargs(
             prompt_adapter_config=self.prompt_adapter_config,
             is_driver_worker=(not self.parallel_config)
             or (rank % self.parallel_config.tensor_parallel_size == 0),
+            observability_config=self.observability_config,
         )
 
     def _get_create_worker_kwargs(

vllm/sequence.py

@@ -91,13 +91,18 @@ class RequestMetrics:
         first_token_time: The time when the first token was generated.
         time_in_queue: The time the request spent in the queue.
         finished_time: The time when the request was finished.
+        scheduler_time: The time spent in the scheduler when this request was being considered by the scheduler.
+        model_forward_time: The time spent in the model forward pass when this request was in the batch.
     """
     arrival_time: float
     last_token_time: float
     first_scheduled_time: Optional[float]
     first_token_time: Optional[float]
     time_in_queue: Optional[float]
     finished_time: Optional[float] = None
+    scheduler_time: Optional[float] = None
+    model_forward_time: Optional[float] = None
+    model_execute_time: Optional[float] = None
 
 
 class SequenceData:
@@ -863,6 +868,13 @@ class SamplerOutput:
     # Optional last hidden states from the model.
     hidden_states: Optional[torch.Tensor] = None
 
+    # Total time spent in the forward pass for this across all workers
+    model_forward_time: Optional[float] = None
+
+    # Total time spent in the model execute function. This will include model forward,
+    #  block/sync across workers, cpu-gpu sync time and sampling time.
+    model_execute_time: Optional[float] = None
+
     def __getitem__(self, idx: int):
         return self.outputs[idx]
 

vllm/tracing.py

@@ -92,7 +92,9 @@ class SpanAttributes(BaseSpanAttributes):
     LLM_LATENCY_TIME_IN_QUEUE = "gen_ai.latency.time_in_queue"
     LLM_LATENCY_TIME_TO_FIRST_TOKEN = "gen_ai.latency.time_to_first_token"
     LLM_LATENCY_E2E = "gen_ai.latency.e2e"
-
+    LLM_LATENCY_TIME_IN_SCHEDULER = "gen_ai.latency.time_in_scheduler"
+    LLM_LATENCY_TIME_IN_MODEL_FORWARD = "gen_ai.latency.time_in_model_forward"
+    LLM_LATENCY_TIME_IN_MODEL_EXECUTE = "gen_ai.latency.time_in_model_execute"
 
 def contains_trace_headers(headers: Mapping[str, str]) -> bool:
     return any(h in headers for h in TRACE_HEADERS)

vllm/worker/model_runner.py

@@ -25,8 +25,8 @@
 
 from vllm.attention import AttentionMetadata, get_attn_backend
 from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig,
-                         ModelConfig, MultiModalConfig, ParallelConfig,
-                         PromptAdapterConfig, SchedulerConfig)
+                         ModelConfig, MultiModalConfig, ObservabilityConfig,
+                         ParallelConfig, PromptAdapterConfig, SchedulerConfig)
 from vllm.distributed import get_pp_group
 from vllm.distributed.parallel_state import graph_capture
 from vllm.inputs import INPUT_REGISTRY
@@ -607,6 +607,7 @@ def __init__(
         prompt_adapter_config: Optional[PromptAdapterConfig] = None,
         multimodal_config: Optional[MultiModalConfig] = None,
         return_hidden_states: bool = False,
+        observability_config: Optional[ObservabilityConfig] = None,
     ):
         self.model_config = model_config
         self.parallel_config = parallel_config
@@ -619,6 +620,7 @@ def __init__(
         self.prompt_adapter_config = prompt_adapter_config
         self.multimodal_config = multimodal_config
         self.return_hidden_states = return_hidden_states
+        self.observability_config = observability_config
 
         self.device = self.device_config.device
         self.pin_memory = is_pin_memory_available()
@@ -1311,6 +1313,9 @@ def execute_model(
             "finished_requests_ids": model_input.finished_requests_ids,
             "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,
         } if self.has_seqlen_agnostic else {}
+        model_forward_start = torch.cuda.Event(enable_timing=True)
+        model_forward_end = torch.cuda.Event(enable_timing=True)
+        model_forward_start.record()
         hidden_or_intermediate_states = model_executable(
             input_ids=model_input.input_tokens,
             positions=model_input.input_positions,
@@ -1319,7 +1324,8 @@ def execute_model(
             intermediate_tensors=intermediate_tensors,
             **multi_modal_kwargs,
             **seqlen_agnostic_kwargs)
-
+        model_forward_end.record()
+
         # Compute the logits in the last pipeline stage.
         if not get_pp_group().is_last_rank:
             return hidden_or_intermediate_states
@@ -1335,6 +1341,13 @@ def execute_model(
             logits=logits,
             sampling_metadata=model_input.sampling_metadata,
         )
+        if self.observability_config.collect_model_forward_time:
+            model_forward_end.synchronize()
+            model_forward_time = model_forward_start.elapsed_time(model_forward_end)    
+            # If there are multiple workers, we are still tracking the latency from the start time
+            # of the driver worker to the end time of the driver worker. The model forward time wil
+            # then end up covering the communication time as well.
+            output.model_forward_time = model_forward_time
 
         if self.return_hidden_states:
             # we only need to pass hidden states of most recent token

vllm/worker/model_runner_base.py

@@ -128,6 +128,13 @@ def build(self, *args, **kwargs) -> T:
         """Build metadata with on-device tensors."""
         raise NotImplementedError
 
+class ModelRunnerOutput:
+    """
+      Model runner output that is used to collect the outputs and stats together.
+    """
+    sampler_output: List[SamplerOutput] = []
+    intermediate_tensors: IntermediateTensors = None
+    forward_time_ms: float = 0
 
 class ModelRunnerBase(ABC, Generic[T]):
     """

vllm/worker/worker.py

@@ -7,8 +7,8 @@
 import torch.distributed
 
 from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig,
-                         ModelConfig, MultiModalConfig, ParallelConfig,
-                         PromptAdapterConfig, SchedulerConfig,
+                         ModelConfig, MultiModalConfig, ObservabilityConfig,
+                         ParallelConfig, PromptAdapterConfig, SchedulerConfig,
                          SpeculativeConfig)
 from vllm.distributed import (ensure_model_parallel_initialized,
                               init_distributed_environment,
@@ -50,6 +50,7 @@ def __init__(
         prompt_adapter_config: Optional[PromptAdapterConfig] = None,
         is_driver_worker: bool = False,
         model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None,
+        observability_config: Optional[ObservabilityConfig] = None,
     ) -> None:
         self.model_config = model_config
         self.parallel_config = parallel_config
@@ -72,6 +73,7 @@ def __init__(
             from vllm.utils import init_cached_hf_modules
             init_cached_hf_modules()
         self.multimodal_config = multimodal_config
+        self.observability_config = observability_config
 
         # Return hidden states from target model if the draft model is an
         # mlp_speculator
@@ -99,6 +101,7 @@ def __init__(
             is_driver_worker=is_driver_worker,
             prompt_adapter_config=prompt_adapter_config,
             multimodal_config=multimodal_config,
+            observability_config=observability_config,
             **speculative_args,
         )
         # Uninitialized cache engine. Will be initialized by

vllm/worker/worker_base.py

@@ -2,6 +2,7 @@
 import importlib
 import os
 from abc import ABC, abstractmethod
+import time
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
 
 import torch
@@ -217,6 +218,7 @@ def execute_model(
         self,
         execute_model_req: Optional[ExecuteModelRequest] = None
     ) -> Optional[List[SamplerOutput]]:
+        start_time = time.time()
         """Executes at least one model step on the given sequences, unless no
         sequences are provided."""
         if self.is_driver_worker:
@@ -273,7 +275,10 @@ def execute_model(
             model_input, self.kv_cache[worker_input.virtual_engine]
             if self.kv_cache is not None else None, intermediate_tensors,
             num_steps)
-
+        end_time = time.time()
+        if output is not None:
+            for o in output:
+                o.model_execute_time =  end_time - start_time
         if not get_pp_group().is_last_rank:
             # output is IntermediateTensors
             get_pp_group().send_tensor_dict(output.tensors)