[Hardware][GPU] Profiler config additional to increase it scope and annotation details (#37524)
Signed-off-by: devalshahamd <deval.shah@amd.com> Signed-off-by: Deval Shah <devashah@amd.com> Signed-off-by: Deval Shah <deval.shah@amd.com> Co-authored-by: Deval Shah <devashah@amd.com>
devalshahamd committed
Jul 17, 2026 at 13:38 UTC
fcd2255d16bd3c62493bae5dee769ce998098f21
4 files changed
+240
-26
tests/v1/worker/test_gpu_profiler.py
+63
-1
@@ -1,10 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from unittest.mock import MagicMock
4
+
5
import pytest
6
5
-from vllm.config import ProfilerConfig
7
+from vllm.config import CUDAGraphMode, ProfilerConfig
8
from vllm.config.profiler import _is_uri_path
9
from vllm.profiler.wrapper import WorkerProfiler
10
+from vllm.v1.core.sched.output import CachedRequestData
11
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
12
+from vllm.v1.worker.gpu_worker import Worker
13
14
15
class ConcreteWorkerProfiler(WorkerProfiler):
@@ -236,3 +241,60 @@ class TestIsUriPath:
241
def test_is_uri_path(self, path, expected):
242
"""Test that _is_uri_path correctly identifies URI vs local paths."""
243
assert _is_uri_path(path) == expected
244
+
245
+
246
+class TestAnnotateProfile:
247
+ """Tests for Worker.annotate_profile() annotation string formatting."""
248
+
249
+ def _annotate(self, detailed: bool) -> str:
250
+ worker = MagicMock()
251
+ worker.vllm_config.profiler_config.detailed_trace_annotation = detailed
252
+ worker.profiler = MagicMock()
253
+
254
+ ctx_req = MagicMock(req_id="ctx1", num_computed_tokens=0)
255
+ cached = CachedRequestData(
256
+ req_ids=["gen1"],
257
+ resumed_req_ids=set(),
258
+ new_token_ids=[],
259
+ all_token_ids={},
260
+ new_block_ids=[],
261
+ num_computed_tokens=[10],
262
+ num_output_tokens=[1],
263
+ )
264
+ sched = MagicMock(
265
+ scheduled_new_reqs=[ctx_req],
266
+ scheduled_cached_reqs=cached,
267
+ num_scheduled_tokens={"ctx1": 4, "gen1": 1},
268
+ )
269
+
270
+ Worker.annotate_profile(worker, sched)
271
+ return worker.profiler.annotate_context_manager.call_args[0][0]
272
+
273
+ def test_simple_format_mixed(self):
274
+ assert self._annotate(detailed=False) == (
275
+ "execute_context_1(4)_generation_1(1)"
276
+ )
277
+
278
+ def test_detailed_format_mixed(self):
279
+ # ctx1: sq=4, sk=4, sqsq=16, sqsk=16 | gen1: sq=1, sk=11, sqsq=1, sqsk=11 | bs=5
280
+ assert self._annotate(detailed=True) == (
281
+ "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)"
282
+ )
283
+
284
+
285
+def test_profiler_entered_during_capture():
286
+ """Profiler is used as a context manager in _warmup_and_capture,
287
+ confirming it is active during the actual graph capture run."""
288
+ runner = MagicMock()
289
+ runner.compilation_config.cudagraph_num_of_warmups = 0
290
+ mock_profiler = MagicMock()
291
+
292
+ GPUModelRunner._warmup_and_capture(
293
+ runner,
294
+ desc=MagicMock(num_tokens=4, uniform=True),
295
+ cudagraph_runtime_mode=CUDAGraphMode.FULL,
296
+ profiler=mock_profiler,
297
+ )
298
+
299
+ mock_profiler.__enter__.assert_called_once()
300
+ mock_profiler.__exit__.assert_called_once()
vllm/config/profiler.py
+17
@@ -66,6 +66,17 @@ class ProfilerConfig:
66
"""If `True`, enables memory profiling in the torch profiler.
67
Disabled by default."""
68
69
+ capture_torch_profiler: bool = False
70
+ """If `True`, enables a torch profiler during CUDA graph capture on rank 0.
71
+ Traces are saved to a `capture_traces` subdirectory under `torch_profiler_dir`.
72
+ Requires `profiler` to be set to 'torch'."""
73
+
74
+ detailed_trace_annotation: bool = False
75
+ """If `True`, uses detailed annotations with roofline metrics (sk, sqsq,
76
+ sqsk) in profiler trace events. If `False`, uses simple annotations with
77
+ only context/generation request counts and token counts.
78
+ Disabled by default."""
79
+
80
ignore_frontend: bool = False
81
"""If `True`, disables the front-end profiling of AsyncLLM when using the
82
'torch' profiler. This is needed to reduce overhead when using delay/limit options,
@@ -144,4 +155,10 @@ class ProfilerConfig:
155
if profiler_dir and not _is_uri_path(profiler_dir):
156
self.torch_profiler_dir = os.path.abspath(os.path.expanduser(profiler_dir))
157
158
+ if self.capture_torch_profiler and self.profiler != "torch":
159
+ raise ValueError(
160
+ "capture_torch_profiler is only applicable when profiler is "
161
+ "set to 'torch'"
162
+ )
163
+
164
return self
vllm/v1/worker/gpu_model_runner.py
+62
-12
@@ -8,7 +8,7 @@ import threading
8
import time
9
from collections import defaultdict
10
from collections.abc import Callable, Iterable, Iterator, Sequence
11
-from contextlib import contextmanager
11
+from contextlib import AbstractContextManager, contextmanager, nullcontext
12
from copy import copy, deepcopy
13
from dataclasses import dataclass, replace
14
from functools import reduce
@@ -6756,6 +6756,44 @@ class GPUModelRunner(
6756
# Capture the large shapes first so that the smaller shapes
6757
# can reuse the memory pool allocated for the large shapes.
6758
set_cudagraph_capturing_enabled(True)
6759
+
6760
+ # Setup torch profiler for graph capture traces (conditional)
6761
+ from vllm.distributed.parallel_state import get_world_group
6762
+
6763
+ local_rank = get_world_group().local_rank
6764
+ enable_profiler = (
6765
+ local_rank == 0
6766
+ ) and self.vllm_config.profiler_config.capture_torch_profiler
6767
+ if enable_profiler:
6768
+ trace_dir = (
6769
+ self.vllm_config.profiler_config.torch_profiler_dir + "/capture_traces"
6770
+ )
6771
+ profiler = torch.profiler.profile(
6772
+ activities=[
6773
+ torch.profiler.ProfilerActivity.CPU,
6774
+ torch.profiler.ProfilerActivity.CUDA,
6775
+ ],
6776
+ record_shapes=True,
6777
+ profile_memory=True,
6778
+ with_stack=True,
6779
+ on_trace_ready=torch.profiler.tensorboard_trace_handler(
6780
+ trace_dir,
6781
+ worker_name=f"graph_capture_rank_{local_rank}",
6782
+ use_gzip=True,
6783
+ ),
6784
+ )
6785
+ logger.info_once(
6786
+ "Rank %d: Torch profiler enabled for CUDA graph capture, "
6787
+ "traces will be saved to: %s",
6788
+ local_rank,
6789
+ trace_dir,
6790
+ )
6791
+ else:
6792
+ profiler = nullcontext()
6793
+ logger.info_once(
6794
+ "Rank %d: Torch profiler disabled for CUDA graph capture", local_rank
6795
+ )
6796
+
6797
with self._freeze_gc(), graph_capture(device=self.device):
6798
torch.accelerator.synchronize()
6799
torch.accelerator.empty_cache()
@@ -6768,6 +6806,7 @@ class GPUModelRunner(
6806
self._capture_cudagraphs(
6807
batch_descriptors=batch_descs,
6808
cudagraph_runtime_mode=runtime_mode,
6809
+ profiler=profiler,
6810
)
6811
torch.accelerator.synchronize()
6812
@@ -6811,7 +6850,10 @@ class GPUModelRunner(
6850
profile_seq_lens: int | None = None,
6851
allow_microbatching: bool = False,
6852
num_warmups: int | None = None,
6853
+ profiler: AbstractContextManager[Any] | None = None,
6854
):
6855
+ if profiler is None:
6856
+ profiler = nullcontext()
6857
if num_warmups is None:
6858
num_warmups = self.compilation_config.cudagraph_num_of_warmups
6859
force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL
@@ -6827,22 +6869,29 @@ class GPUModelRunner(
6869
num_active_loras=desc.num_active_loras,
6870
profile_seq_lens=profile_seq_lens,
6871
)
6830
- self._dummy_run(
6831
- desc.num_tokens,
6832
- cudagraph_runtime_mode=cudagraph_runtime_mode,
6833
- uniform_decode=desc.uniform,
6834
- allow_microbatching=allow_microbatching,
6835
- skip_eplb=True,
6836
- remove_lora=False,
6837
- num_active_loras=desc.num_active_loras,
6838
- is_graph_capturing=True,
6839
- profile_seq_lens=profile_seq_lens,
6840
- )
6872
+ with (
6873
+ profiler,
6874
+ torch.profiler.record_function(
6875
+ f"capture_{desc.num_tokens}_{cudagraph_runtime_mode.name}"
6876
+ ),
6877
+ ):
6878
+ self._dummy_run(
6879
+ desc.num_tokens,
6880
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
6881
+ uniform_decode=desc.uniform,
6882
+ allow_microbatching=allow_microbatching,
6883
+ skip_eplb=True,
6884
+ remove_lora=False,
6885
+ num_active_loras=desc.num_active_loras,
6886
+ is_graph_capturing=True,
6887
+ profile_seq_lens=profile_seq_lens,
6888
+ )
6889
6890
def _capture_cudagraphs(
6891
self,
6892
batch_descriptors: list[BatchDescriptor],
6893
cudagraph_runtime_mode: CUDAGraphMode,
6894
+ profiler: AbstractContextManager[Any] | None = None,
6895
):
6896
assert (
6897
cudagraph_runtime_mode != CUDAGraphMode.NONE
@@ -6885,6 +6934,7 @@ class GPUModelRunner(
6934
batch_desc,
6935
cudagraph_runtime_mode=cudagraph_runtime_mode,
6936
allow_microbatching=allow_microbatching,
6937
+ profiler=profiler,
6938
)
6939
torch.accelerator.synchronize()
6940
self.maybe_remove_all_loras(self.lora_config)
vllm/v1/worker/gpu_worker.py
+98
-13
@@ -975,19 +975,104 @@ class Worker(WorkerBase):
975
976
iteration_details = compute_iteration_details(scheduler_output)
977
978
- annotation = "".join(
979
- [
980
- "execute_context_",
981
- str(iteration_details.num_ctx_requests),
982
- "(",
983
- str(iteration_details.num_ctx_tokens),
984
- ")_generation_",
985
- str(iteration_details.num_generation_requests),
986
- "(",
987
- str(iteration_details.num_generation_tokens),
988
- ")",
989
- ]
990
- )
978
+ if self.vllm_config.profiler_config.detailed_trace_annotation:
979
+ # Compute roofline-model metrics per request, split by phase
980
+ # (context vs generation). These help estimate compute and
981
+ # memory intensity from the trace.
982
+ #
983
+ # Per-request quantities:
984
+ # query_len = number of scheduled (new) tokens for this request
985
+ # seq_len = total sequence length (computed + scheduled tokens)
986
+ #
987
+ # Aggregated across requests in each phase
988
+ # (ctx_=context, gen_=generation):
989
+ # seq_len_sum = sum of seq_len (total KV length)
990
+ # qq_compute = sum of query_len*query_len
991
+ # (proxy for QK^T compute cost)
992
+ # qk_compute = sum of query_len*seq_len
993
+ # (proxy for QK^T compute cost for decode and
994
+ # chunked prefill)
995
+ # total_scheduled_tokens = scheduled tokens across all requests
996
+ ctx_seq_len_sum = 0
997
+ ctx_qq_compute = 0
998
+ ctx_qk_compute = 0
999
+ gen_seq_len_sum = 0
1000
+ gen_qq_compute = 0
1001
+ gen_qk_compute = 0
1002
+ total_scheduled_tokens = 0
1003
+
1004
+ # Build a map of req_id -> num_computed_tokens for all requests
1005
+ new_req_ids = {
1006
+ new_req.req_id for new_req in scheduler_output.scheduled_new_reqs
1007
+ }
1008
+ num_computed_tokens_ids = {
1009
+ new_req.req_id: new_req.num_computed_tokens
1010
+ for new_req in scheduler_output.scheduled_new_reqs
1011
+ }
1012
+ for req_id, num_computed_tokens in zip(
1013
+ scheduler_output.scheduled_cached_reqs.req_ids,
1014
+ scheduler_output.scheduled_cached_reqs.num_computed_tokens,
1015
+ ):
1016
+ num_computed_tokens_ids[req_id] = num_computed_tokens
1017
+
1018
+ # Accumulate per-phase metrics
1019
+ for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items():
1020
+ query_len = num_tokens
1021
+ total_scheduled_tokens += query_len
1022
+ seq_len = num_computed_tokens_ids.get(req_id, 0) + query_len
1023
+ if (
1024
+ scheduler_output.scheduled_cached_reqs.is_context_phase(req_id)
1025
+ or req_id in new_req_ids
1026
+ ):
1027
+ ctx_seq_len_sum += seq_len
1028
+ ctx_qq_compute += query_len * query_len
1029
+ ctx_qk_compute += query_len * seq_len
1030
+ else:
1031
+ gen_seq_len_sum += seq_len
1032
+ gen_qq_compute += query_len * query_len
1033
+ gen_qk_compute += query_len * seq_len
1034
+ annotation = "".join(
1035
+ [
1036
+ "execute_",
1037
+ str(total_scheduled_tokens),
1038
+ "_context_",
1039
+ str(iteration_details.num_ctx_requests),
1040
+ "(sq",
1041
+ str(iteration_details.num_ctx_tokens),
1042
+ "sk",
1043
+ str(ctx_seq_len_sum),
1044
+ "sqsq",
1045
+ str(ctx_qq_compute),
1046
+ "sqsk",
1047
+ str(ctx_qk_compute),
1048
+ ")_generation_",
1049
+ str(iteration_details.num_generation_requests),
1050
+ "(sq",
1051
+ str(iteration_details.num_generation_tokens),
1052
+ "sk",
1053
+ str(gen_seq_len_sum),
1054
+ "sqsq",
1055
+ str(gen_qq_compute),
1056
+ "sqsk",
1057
+ str(gen_qk_compute),
1058
+ ")",
1059
+ ]
1060
+ )
1061
+ else:
1062
+ annotation = "".join(
1063
+ [
1064
+ "execute_context_",
1065
+ str(iteration_details.num_ctx_requests),
1066
+ "(",
1067
+ str(iteration_details.num_ctx_tokens),
1068
+ ")",
1069
+ "_generation_",
1070
+ str(iteration_details.num_generation_requests),
1071
+ "(",
1072
+ str(iteration_details.num_generation_tokens),
1073
+ ")",
1074
+ ]
1075
+ )
1076
return self.profiler.annotate_context_manager(annotation)
1077
1078
@torch.inference_mode()