diff --git a/tests/tracing/test_tracing.py b/tests/tracing/test_tracing.py index 2f8f62cf2..face29f8c 100644 --- a/tests/tracing/test_tracing.py +++ b/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 diff --git a/vllm/config.py b/vllm/config.py index 6f0fdf8bc..fe0487ade 100644 --- a/vllm/config.py +++ b/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 " diff --git a/vllm/core/scheduler.py b/vllm/core/scheduler.py index 6e59c5e0f..44924ef4d 100644 --- a/vllm/core/scheduler.py +++ b/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: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 05bfe7c24..de72f5a8f 100644 --- a/vllm/engine/arg_utils.py +++ b/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 diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py index eabe3b23a..888d0e00c 100644 --- a/vllm/engine/llm_engine.py +++ b/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) diff --git a/vllm/executor/executor_base.py b/vllm/executor/executor_base.py index a848bc709..bc4f54455 100644 --- a/vllm/executor/executor_base.py +++ b/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 diff --git a/vllm/executor/gpu_executor.py b/vllm/executor/gpu_executor.py index 3e77af0e2..57b9e2b33 100644 --- a/vllm/executor/gpu_executor.py +++ b/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( diff --git a/vllm/sequence.py b/vllm/sequence.py index 0cd4c7e71..f516690dd 100644 --- a/vllm/sequence.py +++ b/vllm/sequence.py @@ -91,6 +91,8 @@ 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 @@ -98,6 +100,9 @@ class RequestMetrics: 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] diff --git a/vllm/tracing.py b/vllm/tracing.py index ba6732cab..556ec8623 100644 --- a/vllm/tracing.py +++ b/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) diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py index e63be184a..fc56c75dd 100644 --- a/vllm/worker/model_runner.py +++ b/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 diff --git a/vllm/worker/model_runner_base.py b/vllm/worker/model_runner_base.py index 5fb97025a..bdae08085 100644 --- a/vllm/worker/model_runner_base.py +++ b/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]): """ diff --git a/vllm/worker/worker.py b/vllm/worker/worker.py index f3c379d1a..c88c28d57 100644 --- a/vllm/worker/worker.py +++ b/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 diff --git a/vllm/worker/worker_base.py b/vllm/worker/worker_base.py index 03e3857e2..2a0693e6b 100644 --- a/vllm/worker/worker_base.py +++ b/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)