From fcd2255d16bd3c62493bae5dee769ce998098f21 Mon Sep 17 00:00:00 2001 From: devalshahamd Date: Fri, 17 Jul 2026 13:38:59 -0700 Subject: [PATCH] [Hardware][GPU] Profiler config additional to increase it scope and annotation details (#37524) Signed-off-by: devalshahamd Signed-off-by: Deval Shah Signed-off-by: Deval Shah Co-authored-by: Deval Shah --- tests/v1/worker/test_gpu_profiler.py | 64 ++++++++++++++- vllm/config/profiler.py | 17 ++++ vllm/v1/worker/gpu_model_runner.py | 74 +++++++++++++++--- vllm/v1/worker/gpu_worker.py | 111 +++++++++++++++++++++++---- 4 files changed, 240 insertions(+), 26 deletions(-) diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index ca22f3c9da6..8ff89354ea9 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -1,10 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from unittest.mock import MagicMock + import pytest -from vllm.config import ProfilerConfig +from vllm.config import CUDAGraphMode, ProfilerConfig from vllm.config.profiler import _is_uri_path from vllm.profiler.wrapper import WorkerProfiler +from vllm.v1.core.sched.output import CachedRequestData +from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm.v1.worker.gpu_worker import Worker class ConcreteWorkerProfiler(WorkerProfiler): @@ -236,3 +241,60 @@ class TestIsUriPath: def test_is_uri_path(self, path, expected): """Test that _is_uri_path correctly identifies URI vs local paths.""" assert _is_uri_path(path) == expected + + +class TestAnnotateProfile: + """Tests for Worker.annotate_profile() annotation string formatting.""" + + def _annotate(self, detailed: bool) -> str: + worker = MagicMock() + worker.vllm_config.profiler_config.detailed_trace_annotation = detailed + worker.profiler = MagicMock() + + ctx_req = MagicMock(req_id="ctx1", num_computed_tokens=0) + cached = CachedRequestData( + req_ids=["gen1"], + resumed_req_ids=set(), + new_token_ids=[], + all_token_ids={}, + new_block_ids=[], + num_computed_tokens=[10], + num_output_tokens=[1], + ) + sched = MagicMock( + scheduled_new_reqs=[ctx_req], + scheduled_cached_reqs=cached, + num_scheduled_tokens={"ctx1": 4, "gen1": 1}, + ) + + Worker.annotate_profile(worker, sched) + return worker.profiler.annotate_context_manager.call_args[0][0] + + def test_simple_format_mixed(self): + assert self._annotate(detailed=False) == ( + "execute_context_1(4)_generation_1(1)" + ) + + def test_detailed_format_mixed(self): + # ctx1: sq=4, sk=4, sqsq=16, sqsk=16 | gen1: sq=1, sk=11, sqsq=1, sqsk=11 | bs=5 + assert self._annotate(detailed=True) == ( + "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)" + ) + + +def test_profiler_entered_during_capture(): + """Profiler is used as a context manager in _warmup_and_capture, + confirming it is active during the actual graph capture run.""" + runner = MagicMock() + runner.compilation_config.cudagraph_num_of_warmups = 0 + mock_profiler = MagicMock() + + GPUModelRunner._warmup_and_capture( + runner, + desc=MagicMock(num_tokens=4, uniform=True), + cudagraph_runtime_mode=CUDAGraphMode.FULL, + profiler=mock_profiler, + ) + + mock_profiler.__enter__.assert_called_once() + mock_profiler.__exit__.assert_called_once() diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index 68fa78854b4..f0e29d08f38 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -66,6 +66,17 @@ class ProfilerConfig: """If `True`, enables memory profiling in the torch profiler. Disabled by default.""" + capture_torch_profiler: bool = False + """If `True`, enables a torch profiler during CUDA graph capture on rank 0. + Traces are saved to a `capture_traces` subdirectory under `torch_profiler_dir`. + Requires `profiler` to be set to 'torch'.""" + + detailed_trace_annotation: bool = False + """If `True`, uses detailed annotations with roofline metrics (sk, sqsq, + sqsk) in profiler trace events. If `False`, uses simple annotations with + only context/generation request counts and token counts. + Disabled by default.""" + ignore_frontend: bool = False """If `True`, disables the front-end profiling of AsyncLLM when using the 'torch' profiler. This is needed to reduce overhead when using delay/limit options, @@ -144,4 +155,10 @@ class ProfilerConfig: if profiler_dir and not _is_uri_path(profiler_dir): self.torch_profiler_dir = os.path.abspath(os.path.expanduser(profiler_dir)) + if self.capture_torch_profiler and self.profiler != "torch": + raise ValueError( + "capture_torch_profiler is only applicable when profiler is " + "set to 'torch'" + ) + return self diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 79a905adada..2c7adaaf2ae 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -8,7 +8,7 @@ import threading import time from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager, nullcontext from copy import copy, deepcopy from dataclasses import dataclass, replace from functools import reduce @@ -6756,6 +6756,44 @@ class GPUModelRunner( # Capture the large shapes first so that the smaller shapes # can reuse the memory pool allocated for the large shapes. set_cudagraph_capturing_enabled(True) + + # Setup torch profiler for graph capture traces (conditional) + from vllm.distributed.parallel_state import get_world_group + + local_rank = get_world_group().local_rank + enable_profiler = ( + local_rank == 0 + ) and self.vllm_config.profiler_config.capture_torch_profiler + if enable_profiler: + trace_dir = ( + self.vllm_config.profiler_config.torch_profiler_dir + "/capture_traces" + ) + profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + profile_memory=True, + with_stack=True, + on_trace_ready=torch.profiler.tensorboard_trace_handler( + trace_dir, + worker_name=f"graph_capture_rank_{local_rank}", + use_gzip=True, + ), + ) + logger.info_once( + "Rank %d: Torch profiler enabled for CUDA graph capture, " + "traces will be saved to: %s", + local_rank, + trace_dir, + ) + else: + profiler = nullcontext() + logger.info_once( + "Rank %d: Torch profiler disabled for CUDA graph capture", local_rank + ) + with self._freeze_gc(), graph_capture(device=self.device): torch.accelerator.synchronize() torch.accelerator.empty_cache() @@ -6768,6 +6806,7 @@ class GPUModelRunner( self._capture_cudagraphs( batch_descriptors=batch_descs, cudagraph_runtime_mode=runtime_mode, + profiler=profiler, ) torch.accelerator.synchronize() @@ -6811,7 +6850,10 @@ class GPUModelRunner( profile_seq_lens: int | None = None, allow_microbatching: bool = False, num_warmups: int | None = None, + profiler: AbstractContextManager[Any] | None = None, ): + if profiler is None: + profiler = nullcontext() if num_warmups is None: num_warmups = self.compilation_config.cudagraph_num_of_warmups force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL @@ -6827,22 +6869,29 @@ class GPUModelRunner( num_active_loras=desc.num_active_loras, profile_seq_lens=profile_seq_lens, ) - self._dummy_run( - desc.num_tokens, - cudagraph_runtime_mode=cudagraph_runtime_mode, - uniform_decode=desc.uniform, - allow_microbatching=allow_microbatching, - skip_eplb=True, - remove_lora=False, - num_active_loras=desc.num_active_loras, - is_graph_capturing=True, - profile_seq_lens=profile_seq_lens, - ) + with ( + profiler, + torch.profiler.record_function( + f"capture_{desc.num_tokens}_{cudagraph_runtime_mode.name}" + ), + ): + self._dummy_run( + desc.num_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + uniform_decode=desc.uniform, + allow_microbatching=allow_microbatching, + skip_eplb=True, + remove_lora=False, + num_active_loras=desc.num_active_loras, + is_graph_capturing=True, + profile_seq_lens=profile_seq_lens, + ) def _capture_cudagraphs( self, batch_descriptors: list[BatchDescriptor], cudagraph_runtime_mode: CUDAGraphMode, + profiler: AbstractContextManager[Any] | None = None, ): assert ( cudagraph_runtime_mode != CUDAGraphMode.NONE @@ -6885,6 +6934,7 @@ class GPUModelRunner( batch_desc, cudagraph_runtime_mode=cudagraph_runtime_mode, allow_microbatching=allow_microbatching, + profiler=profiler, ) torch.accelerator.synchronize() self.maybe_remove_all_loras(self.lora_config) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 5fb0c387737..9c20df0d18c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -975,19 +975,104 @@ class Worker(WorkerBase): iteration_details = compute_iteration_details(scheduler_output) - annotation = "".join( - [ - "execute_context_", - str(iteration_details.num_ctx_requests), - "(", - str(iteration_details.num_ctx_tokens), - ")_generation_", - str(iteration_details.num_generation_requests), - "(", - str(iteration_details.num_generation_tokens), - ")", - ] - ) + if self.vllm_config.profiler_config.detailed_trace_annotation: + # Compute roofline-model metrics per request, split by phase + # (context vs generation). These help estimate compute and + # memory intensity from the trace. + # + # Per-request quantities: + # query_len = number of scheduled (new) tokens for this request + # seq_len = total sequence length (computed + scheduled tokens) + # + # Aggregated across requests in each phase + # (ctx_=context, gen_=generation): + # seq_len_sum = sum of seq_len (total KV length) + # qq_compute = sum of query_len*query_len + # (proxy for QK^T compute cost) + # qk_compute = sum of query_len*seq_len + # (proxy for QK^T compute cost for decode and + # chunked prefill) + # total_scheduled_tokens = scheduled tokens across all requests + ctx_seq_len_sum = 0 + ctx_qq_compute = 0 + ctx_qk_compute = 0 + gen_seq_len_sum = 0 + gen_qq_compute = 0 + gen_qk_compute = 0 + total_scheduled_tokens = 0 + + # Build a map of req_id -> num_computed_tokens for all requests + new_req_ids = { + new_req.req_id for new_req in scheduler_output.scheduled_new_reqs + } + num_computed_tokens_ids = { + new_req.req_id: new_req.num_computed_tokens + for new_req in scheduler_output.scheduled_new_reqs + } + for req_id, num_computed_tokens in zip( + scheduler_output.scheduled_cached_reqs.req_ids, + scheduler_output.scheduled_cached_reqs.num_computed_tokens, + ): + num_computed_tokens_ids[req_id] = num_computed_tokens + + # Accumulate per-phase metrics + for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items(): + query_len = num_tokens + total_scheduled_tokens += query_len + seq_len = num_computed_tokens_ids.get(req_id, 0) + query_len + if ( + scheduler_output.scheduled_cached_reqs.is_context_phase(req_id) + or req_id in new_req_ids + ): + ctx_seq_len_sum += seq_len + ctx_qq_compute += query_len * query_len + ctx_qk_compute += query_len * seq_len + else: + gen_seq_len_sum += seq_len + gen_qq_compute += query_len * query_len + gen_qk_compute += query_len * seq_len + annotation = "".join( + [ + "execute_", + str(total_scheduled_tokens), + "_context_", + str(iteration_details.num_ctx_requests), + "(sq", + str(iteration_details.num_ctx_tokens), + "sk", + str(ctx_seq_len_sum), + "sqsq", + str(ctx_qq_compute), + "sqsk", + str(ctx_qk_compute), + ")_generation_", + str(iteration_details.num_generation_requests), + "(sq", + str(iteration_details.num_generation_tokens), + "sk", + str(gen_seq_len_sum), + "sqsq", + str(gen_qq_compute), + "sqsk", + str(gen_qk_compute), + ")", + ] + ) + else: + annotation = "".join( + [ + "execute_context_", + str(iteration_details.num_ctx_requests), + "(", + str(iteration_details.num_ctx_tokens), + ")", + "_generation_", + str(iteration_details.num_generation_requests), + "(", + str(iteration_details.num_generation_tokens), + ")", + ] + ) return self.profiler.annotate_context_manager(annotation) @torch.inference_mode()