@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e858bfe051
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 11, 18:45 UTC

[Cleanup] Refactor profiling env vars into a CLI config (#29912)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Benjamin Chislett committed Dec 9, 2025 at 13:29 UTC e858bfe05167a3bbb064e283da5a1a7709dee24e
22 files changed +437 -256
benchmarks/auto_tune/auto_tune.sh
+3 -2
@@ -96,8 +96,9 @@ start_server() {
96 # This correctly passes each element as a separate argument.
97 if [[ -n "$profile_dir" ]]; then
98 # Start server with profiling enabled
99 - VLLM_SERVER_DEV_MODE=1 VLLM_TORCH_PROFILER_DIR=$profile_dir \
100 - vllm serve "${common_args_array[@]}" > "$vllm_log" 2>&1 &
99 + local profile_config_json="{\"profiler\": \"torch\", \"torch_profiler_dir\": \"$profile_dir\"}"
100 + VLLM_SERVER_DEV_MODE=1 \
101 + vllm serve --profiler-config "$profile_config_json" "${common_args_array[@]}" > "$vllm_log" 2>&1 &
102 else
103 # Start server without profiling
104 VLLM_SERVER_DEV_MODE=1 \
benchmarks/benchmark_serving_structured_output.py
+1 -2
@@ -963,8 +963,7 @@ def create_argument_parser():
963 parser.add_argument(
964 "--profile",
965 action="store_true",
966 - help="Use Torch Profiler. The endpoint must be launched with "
967 - "VLLM_TORCH_PROFILER_DIR to enable profiler.",
966 + help="Use vLLM Profiling. --profiler-config must be provided on the server.",
967 )
968 parser.add_argument(
969 "--result-dir",
docs/api/README.md
+1
@@ -15,6 +15,7 @@ API documentation for vLLM's configuration classes.
15 - [vllm.config.MultiModalConfig][]
16 - [vllm.config.PoolerConfig][]
17 - [vllm.config.StructuredOutputsConfig][]
18 +- [vllm.config.ProfilerConfig][]
19 - [vllm.config.ObservabilityConfig][]
20 - [vllm.config.KVTransferConfig][]
21 - [vllm.config.CompilationConfig][]
docs/contributing/profiling.md
+10 -13
@@ -5,16 +5,15 @@
5
6 ## Profile with PyTorch Profiler
7
8 -We support tracing vLLM workers using the `torch.profiler` module. You can enable tracing by setting the `VLLM_TORCH_PROFILER_DIR` environment variable to the directory where you want to save the traces: `VLLM_TORCH_PROFILER_DIR=/mnt/traces/`. Additionally, you can control the profiling content by specifying the following environment variables:
8 +We support tracing vLLM workers using the `torch.profiler` module. You can enable the torch profiler by setting `--profiler-config`
9 +when launching the server, and setting the entries `profiler` to `'torch'` and `torch_profiler_dir` to the directory where you want to save the traces. Additionally, you can control the profiling content by specifying the following additional arguments in the config:
10
10 -- `VLLM_TORCH_PROFILER_RECORD_SHAPES=1` to enable recording Tensor Shapes, off by default
11 -- `VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY=1` to record memory, off by default
12 -- `VLLM_TORCH_PROFILER_WITH_STACK=1` to enable recording stack information, on by default
13 -- `VLLM_TORCH_PROFILER_WITH_FLOPS=1` to enable recording FLOPs, off by default
14 -- `VLLM_TORCH_PROFILER_USE_GZIP=0` to disable gzip-compressing profiling files, on by default
15 -- `VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL=0` to disable dumping and printing the aggregated CUDA self time table, on by default
16 -
17 -The OpenAI server also needs to be started with the `VLLM_TORCH_PROFILER_DIR` environment variable set.
11 +- `torch_profiler_record_shapes` to enable recording Tensor Shapes, off by default
12 +- `torch_profiler_with_memory` to record memory, off by default
13 +- `torch_profiler_with_stack` to enable recording stack information, on by default
14 +- `torch_profiler_with_flops` to enable recording FLOPs, off by default
15 +- `torch_profiler_use_gzip` to control gzip-compressing profiling files, on by default
16 +- `torch_profiler_dump_cuda_time_total` to control dumping and printing the aggregated CUDA self time table, on by default
17
18 When using `vllm bench serve`, you can enable profiling by passing the `--profile` flag.
19
@@ -40,8 +39,7 @@ Refer to [examples/offline_inference/simple_profiling.py](../../examples/offline
39 #### OpenAI Server
40
41 ```bash
43 -VLLM_TORCH_PROFILER_DIR=./vllm_profile \
44 - vllm serve meta-llama/Llama-3.1-8B-Instruct
42 +vllm serve meta-llama/Llama-3.1-8B-Instruct --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
43 ```
44
45 vllm bench command:
@@ -104,13 +102,12 @@ To profile the server, you will want to prepend your `vllm serve` command with `
102
103 ```bash
104 # server
107 -VLLM_TORCH_CUDA_PROFILE=1 \
105 nsys profile \
106 --trace-fork-before-exec=true \
107 --cuda-graph-trace=node \
108 --capture-range=cudaProfilerApi \
109 --capture-range-end repeat \
113 - vllm serve meta-llama/Llama-3.1-8B-Instruct
110 + vllm serve meta-llama/Llama-3.1-8B-Instruct --profiler-config.profiler cuda
111
112 # client
113 vllm bench serve \
examples/offline_inference/simple_profiling.py
+8 -5
@@ -1,14 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import os
4 import time
5
6 from vllm import LLM, SamplingParams
7
9 -# enable torch profiler, can also be set on cmd line
10 -os.environ["VLLM_TORCH_PROFILER_DIR"] = "./vllm_profile"
11 -
8 # Sample prompts.
9 prompts = [
10 "Hello, my name is",
@@ -22,7 +18,14 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
18
19 def main():
20 # Create an LLM.
25 - llm = LLM(model="facebook/opt-125m", tensor_parallel_size=1)
21 + llm = LLM(
22 + model="facebook/opt-125m",
23 + tensor_parallel_size=1,
24 + profiler_config={
25 + "profiler": "torch",
26 + "torch_profiler_dir": "./vllm_profile",
27 + },
28 + )
29
30 llm.start_profile()
31
tests/v1/worker/test_gpu_profiler.py
+36 -35
@@ -2,8 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4
5 -import vllm.envs as envs
6 -from vllm.profiler.gpu_profiler import WorkerProfiler
5 +from vllm.config import ProfilerConfig
6 +from vllm.profiler.wrapper import WorkerProfiler
7
8
9 class ConcreteWorkerProfiler(WorkerProfiler):
@@ -11,11 +11,11 @@ class ConcreteWorkerProfiler(WorkerProfiler):
11 A basic implementation of a worker profiler for testing purposes.
12 """
13
14 - def __init__(self):
14 + def __init__(self, profiler_config: ProfilerConfig):
15 self.start_call_count = 0
16 self.stop_call_count = 0
17 self.should_fail_start = False
18 - super().__init__()
18 + super().__init__(profiler_config)
19
20 def _start(self) -> None:
21 if self.should_fail_start:
@@ -26,17 +26,19 @@ class ConcreteWorkerProfiler(WorkerProfiler):
26 self.stop_call_count += 1
27
28
29 -@pytest.fixture(autouse=True)
30 -def reset_mocks():
31 - """Fixture to reset mocks and env variables before each test."""
32 - envs.VLLM_PROFILER_DELAY_ITERS = 0
33 - envs.VLLM_PROFILER_MAX_ITERS = 0
29 +@pytest.fixture
30 +def default_profiler_config():
31 + return ProfilerConfig(
32 + profiler="torch",
33 + torch_profiler_dir="/tmp/mock",
34 + delay_iterations=0,
35 + max_iterations=0,
36 + )
37
38
36 -def test_immediate_start_stop():
39 +def test_immediate_start_stop(default_profiler_config):
40 """Test standard start without delay."""
38 - profiler = ConcreteWorkerProfiler()
39 -
41 + profiler = ConcreteWorkerProfiler(default_profiler_config)
42 profiler.start()
43 assert profiler._running is True
44 assert profiler._active is True
@@ -48,10 +50,10 @@ def test_immediate_start_stop():
50 assert profiler.stop_call_count == 1
51
52
51 -def test_delayed_start():
53 +def test_delayed_start(default_profiler_config):
54 """Test that profiler waits for N steps before actually starting."""
53 - envs.VLLM_PROFILER_DELAY_ITERS = 2
54 - profiler = ConcreteWorkerProfiler()
55 + default_profiler_config.delay_iterations = 2
56 + profiler = ConcreteWorkerProfiler(default_profiler_config)
57
58 # User requests start
59 profiler.start()
@@ -71,10 +73,10 @@ def test_delayed_start():
73 assert profiler.start_call_count == 1
74
75
74 -def test_max_iterations():
76 +def test_max_iterations(default_profiler_config):
77 """Test that profiler stops automatically after max iterations."""
76 - envs.VLLM_PROFILER_MAX_ITERS = 2
77 - profiler = ConcreteWorkerProfiler()
78 + default_profiler_config.max_iterations = 2
79 + profiler = ConcreteWorkerProfiler(default_profiler_config)
80
81 profiler.start()
82 assert profiler._running is True
@@ -95,12 +97,11 @@ def test_max_iterations():
97 assert profiler.stop_call_count == 1
98
99
98 -def test_delayed_start_and_max_iters():
100 +def test_delayed_start_and_max_iters(default_profiler_config):
101 """Test combined delayed start and max iterations."""
100 - envs.VLLM_PROFILER_DELAY_ITERS = 2
101 - envs.VLLM_PROFILER_MAX_ITERS = 2
102 - profiler = ConcreteWorkerProfiler()
103 -
102 + default_profiler_config.delay_iterations = 2
103 + default_profiler_config.max_iterations = 2
104 + profiler = ConcreteWorkerProfiler(default_profiler_config)
105 profiler.start()
106
107 # Step 1
@@ -127,9 +128,9 @@ def test_delayed_start_and_max_iters():
128 assert profiler.stop_call_count == 1
129
130
130 -def test_idempotency():
131 +def test_idempotency(default_profiler_config):
132 """Test that calling start/stop multiple times doesn't break logic."""
132 - profiler = ConcreteWorkerProfiler()
133 + profiler = ConcreteWorkerProfiler(default_profiler_config)
134
135 # Double Start
136 profiler.start()
@@ -142,10 +143,10 @@ def test_idempotency():
143 assert profiler.stop_call_count == 1 # Should only stop once
144
145
145 -def test_step_inactive():
146 +def test_step_inactive(default_profiler_config):
147 """Test that stepping while inactive does nothing."""
147 - envs.VLLM_PROFILER_DELAY_ITERS = 2
148 - profiler = ConcreteWorkerProfiler()
148 + default_profiler_config.delay_iterations = 2
149 + profiler = ConcreteWorkerProfiler(default_profiler_config)
150
151 # Not started yet
152 profiler.step()
@@ -155,9 +156,9 @@ def test_step_inactive():
156 assert profiler.start_call_count == 0
157
158
158 -def test_start_failure():
159 +def test_start_failure(default_profiler_config):
160 """Test behavior when the underlying _start method raises exception."""
160 - profiler = ConcreteWorkerProfiler()
161 + profiler = ConcreteWorkerProfiler(default_profiler_config)
162 profiler.should_fail_start = True
163
164 profiler.start()
@@ -168,9 +169,9 @@ def test_start_failure():
169 assert profiler.start_call_count == 0 # Logic failed inside start
170
171
171 -def test_shutdown():
172 +def test_shutdown(default_profiler_config):
173 """Test that shutdown calls stop only if running."""
173 - profiler = ConcreteWorkerProfiler()
174 + profiler = ConcreteWorkerProfiler(default_profiler_config)
175
176 # Case 1: Not running
177 profiler.shutdown()
@@ -182,10 +183,10 @@ def test_shutdown():
183 assert profiler.stop_call_count == 1
184
185
185 -def test_mixed_delay_and_stop():
186 +def test_mixed_delay_and_stop(default_profiler_config):
187 """Test manual stop during the delay period."""
187 - envs.VLLM_PROFILER_DELAY_ITERS = 5
188 - profiler = ConcreteWorkerProfiler()
188 + default_profiler_config.delay_iterations = 5
189 + profiler = ConcreteWorkerProfiler(default_profiler_config)
190
191 profiler.start()
192 profiler.step()
vllm/benchmarks/latency.py
+5 -7
@@ -12,7 +12,6 @@ from typing import Any
12 import numpy as np
13 from tqdm import tqdm
14
15 -import vllm.envs as envs
15 from vllm.benchmarks.lib.utils import convert_to_pytorch_benchmark_format, write_to_json
16 from vllm.engine.arg_utils import EngineArgs
17 from vllm.inputs import PromptType
@@ -79,12 +78,11 @@ def add_cli_args(parser: argparse.ArgumentParser):
78
79
80 def main(args: argparse.Namespace):
82 - if args.profile and not envs.VLLM_TORCH_PROFILER_DIR:
83 - raise OSError(
84 - "The environment variable 'VLLM_TORCH_PROFILER_DIR' is not set. "
85 - "Please set it to a valid path to use torch profiler."
86 - )
81 engine_args = EngineArgs.from_cli_args(args)
82 + if args.profile and not engine_args.profiler_config.profiler == "torch":
83 + raise ValueError(
84 + "The torch profiler is not enabled. Please provide profiler_config."
85 + )
86
87 # Lazy import to avoid importing LLM when the bench command is not selected.
88 from vllm import LLM, SamplingParams
@@ -144,7 +142,7 @@ def main(args: argparse.Namespace):
142 run_to_completion(profile_dir=None)
143
144 if args.profile:
147 - profile_dir = envs.VLLM_TORCH_PROFILER_DIR
145 + profile_dir = engine_args.profiler_config.torch_profiler_dir
146 print(f"Profiling (results will be saved to '{profile_dir}')...")
147 run_to_completion(profile_dir=profile_dir)
148 return
vllm/benchmarks/serve.py
+1 -2
@@ -1097,8 +1097,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
1097 parser.add_argument(
1098 "--profile",
1099 action="store_true",
1100 - help="Use Torch Profiler. The endpoint must be launched with "
1101 - "VLLM_TORCH_PROFILER_DIR to enable profiler.",
1100 + help="Use vLLM Profiling. --profiler-config must be provided on the server.",
1101 )
1102 parser.add_argument(
1103 "--save-result",
vllm/benchmarks/throughput.py
+1 -2
@@ -655,8 +655,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
655 "--profile",
656 action="store_true",
657 default=False,
658 - help="Use Torch Profiler. The env variable "
659 - "VLLM_TORCH_PROFILER_DIR must be set to enable profiler.",
658 + help="Use vLLM Profiling. --profiler-config must be provided on the server.",
659 )
660
661 # prefix repetition dataset
vllm/config/__init__.py
+3
@@ -24,6 +24,7 @@ from vllm.config.multimodal import MultiModalConfig
24 from vllm.config.observability import ObservabilityConfig
25 from vllm.config.parallel import EPLBConfig, ParallelConfig
26 from vllm.config.pooler import PoolerConfig
27 +from vllm.config.profiler import ProfilerConfig
28 from vllm.config.scheduler import SchedulerConfig
29 from vllm.config.speculative import SpeculativeConfig
30 from vllm.config.speech_to_text import SpeechToTextConfig
@@ -89,6 +90,8 @@ __all__ = [
90 "SpeechToTextConfig",
91 # From vllm.config.structured_outputs
92 "StructuredOutputsConfig",
93 + # From vllm.config.profiler
94 + "ProfilerConfig",
95 # From vllm.config.utils
96 "ConfigType",
97 "SupportsMetricsInfo",
vllm/config/profiler.py new
+199
@@ -0,0 +1,199 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import os
5 +from typing import Any, Literal
6 +
7 +from pydantic import Field, model_validator
8 +from pydantic.dataclasses import dataclass
9 +from typing_extensions import Self
10 +
11 +import vllm.envs as envs
12 +from vllm.config.utils import config
13 +from vllm.logger import init_logger
14 +from vllm.utils.hashing import safe_hash
15 +
16 +logger = init_logger(__name__)
17 +
18 +ProfilerKind = Literal["torch", "cuda"]
19 +
20 +
21 +@config
22 +@dataclass
23 +class ProfilerConfig:
24 + """Dataclass which contains profiler config for the engine."""
25 +
26 + profiler: ProfilerKind | None = None
27 + """Which profiler to use. Defaults to None. Options are:
28 +
29 + - 'torch': Use PyTorch profiler.\n
30 + - 'cuda': Use CUDA profiler."""
31 +
32 + torch_profiler_dir: str = ""
33 + """Directory to save torch profiler traces. Both AsyncLLM's CPU traces and
34 + worker's traces (CPU & GPU) will be saved under this directory. Note that
35 + it must be an absolute path."""
36 +
37 + torch_profiler_with_stack: bool = True
38 + """If `True`, enables stack tracing in the torch profiler. Enabled by default."""
39 +
40 + torch_profiler_with_flops: bool = False
41 + """If `True`, enables FLOPS counting in the torch profiler. Disabled by default."""
42 +
43 + torch_profiler_use_gzip: bool = True
44 + """If `True`, saves torch profiler traces in gzip format. Enabled by default"""
45 +
46 + torch_profiler_dump_cuda_time_total: bool = True
47 + """If `True`, dumps total CUDA time in torch profiler traces. Enabled by default."""
48 +
49 + torch_profiler_record_shapes: bool = False
50 + """If `True`, records tensor shapes in the torch profiler. Disabled by default."""
51 +
52 + torch_profiler_with_memory: bool = False
53 + """If `True`, enables memory profiling in the torch profiler.
54 + Disabled by default."""
55 +
56 + ignore_frontend: bool = False
57 + """If `True`, disables the front-end profiling of AsyncLLM when using the
58 + 'torch' profiler. This is needed to reduce overhead when using delay/limit options,
59 + since the front-end profiling does not track iterations and will capture the
60 + entire range.
61 + """
62 +
63 + delay_iterations: int = Field(default=0, ge=0)
64 + """Number of engine iterations to skip before starting profiling.
65 + Defaults to 0, meaning profiling starts immediately after receiving /start_profile.
66 + """
67 +
68 + max_iterations: int = Field(default=0, ge=0)
69 + """Maximum number of engine iterations to profile after starting profiling.
70 + Defaults to 0, meaning no limit.
71 + """
72 +
73 + def compute_hash(self) -> str:
74 + """
75 + WARNING: Whenever a new field is added to this config,
76 + ensure that it is included in the factors list if
77 + it affects the computation graph.
78 +
79 + Provide a hash that uniquely identifies all the configs
80 + that affect the structure of the computation
81 + graph from input ids/embeddings to the final hidden states,
82 + excluding anything before input ids/embeddings and after
83 + the final hidden states.
84 + """
85 + # no factors to consider.
86 + # this config will not affect the computation graph.
87 + factors: list[Any] = []
88 + hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
89 + return hash_str
90 +
91 + def _get_from_env_if_set(self, field_name: str, env_var_name: str) -> None:
92 + """Get field from env var if set, with deprecation warning."""
93 +
94 + if envs.is_set(env_var_name):
95 + value = getattr(envs, env_var_name)
96 + logger.warning_once(
97 + "Using %s environment variable is deprecated and will be removed in "
98 + "v0.14.0 or v1.0.0, whichever is soonest. Please use "
99 + "--profiler-config.%s command line argument or "
100 + "ProfilerConfig(%s=...) config field instead.",
101 + env_var_name,
102 + field_name,
103 + field_name,
104 + )
105 + return value
106 + return None
107 +
108 + def _set_from_env_if_set(
109 + self,
110 + field_name: str,
111 + env_var_name: str,
112 + to_bool: bool = True,
113 + to_int: bool = False,
114 + ) -> None:
115 + """Set field from env var if set, with deprecation warning."""
116 + value = self._get_from_env_if_set(field_name, env_var_name)
117 + if value is not None:
118 + if to_bool:
119 + value = value == "1"
120 + if to_int:
121 + value = int(value)
122 + setattr(self, field_name, value)
123 +
124 + @model_validator(mode="after")
125 + def _validate_profiler_config(self) -> Self:
126 + maybe_use_cuda_profiler = self._get_from_env_if_set(
127 + "profiler", "VLLM_TORCH_CUDA_PROFILE"
128 + )
129 + if maybe_use_cuda_profiler is not None:
130 + self.profiler = "cuda" if maybe_use_cuda_profiler == "1" else None
131 + else:
132 + self._set_from_env_if_set(
133 + "torch_profiler_dir", "VLLM_TORCH_PROFILER_DIR", to_bool=False
134 + )
135 + if self.torch_profiler_dir:
136 + self.profiler = "torch"
137 + self._set_from_env_if_set(
138 + "torch_profiler_record_shapes",
139 + "VLLM_TORCH_PROFILER_RECORD_SHAPES",
140 + )
141 + self._set_from_env_if_set(
142 + "torch_profiler_with_memory",
143 + "VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY",
144 + )
145 + self._set_from_env_if_set(
146 + "torch_profiler_with_stack",
147 + "VLLM_TORCH_PROFILER_WITH_STACK",
148 + )
149 + self._set_from_env_if_set(
150 + "torch_profiler_with_flops",
151 + "VLLM_TORCH_PROFILER_WITH_FLOPS",
152 + )
153 + self._set_from_env_if_set(
154 + "ignore_frontend",
155 + "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM",
156 + )
157 + self._set_from_env_if_set(
158 + "torch_profiler_use_gzip",
159 + "VLLM_TORCH_PROFILER_USE_GZIP",
160 + )
161 + self._set_from_env_if_set(
162 + "torch_profiler_dump_cuda_time_total",
163 + "VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL",
164 + )
165 +
166 + self._set_from_env_if_set(
167 + "delay_iterations", "VLLM_PROFILER_DELAY_ITERS", to_bool=False, to_int=True
168 + )
169 + self._set_from_env_if_set(
170 + "max_iterations", "VLLM_PROFILER_MAX_ITERS", to_bool=False, to_int=True
171 + )
172 +
173 + has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0
174 + if self.profiler == "torch" and has_delay_or_limit and not self.ignore_frontend:
175 + logger.warning_once(
176 + "Using 'torch' profiler with delay_iterations or max_iterations "
177 + "while ignore_frontend is False may result in high overhead."
178 + )
179 +
180 + profiler_dir = self.torch_profiler_dir
181 + if profiler_dir and self.profiler != "torch":
182 + raise ValueError(
183 + "torch_profiler_dir is only applicable when profiler is set to 'torch'"
184 + )
185 + if self.profiler == "torch" and not profiler_dir:
186 + raise ValueError("torch_profiler_dir must be set when profiler is 'torch'")
187 +
188 + if profiler_dir:
189 + is_gs_path = (
190 + profiler_dir.startswith("gs://")
191 + and profiler_dir[5:]
192 + and profiler_dir[5] != "/"
193 + )
194 + if not is_gs_path:
195 + self.torch_profiler_dir = os.path.abspath(
196 + os.path.expanduser(profiler_dir)
197 + )
198 +
199 + return self
vllm/config/vllm.py
+5
@@ -39,6 +39,7 @@ from .lora import LoRAConfig
39 from .model import ModelConfig
40 from .observability import ObservabilityConfig
41 from .parallel import ParallelConfig
42 +from .profiler import ProfilerConfig
43 from .scheduler import SchedulerConfig
44 from .speculative import SpeculativeConfig
45 from .structured_outputs import StructuredOutputsConfig
@@ -218,6 +219,8 @@ class VllmConfig:
219 You can specify the full compilation config like so:
220 `{"mode": 3, "cudagraph_capture_sizes": [1, 2, 4, 8]}`
221 """
222 + profiler_config: ProfilerConfig = Field(default_factory=ProfilerConfig)
223 + """Profiling configuration."""
224 kv_transfer_config: KVTransferConfig | None = None
225 """The configurations for distributed KV cache transfer."""
226 kv_events_config: KVEventsConfig | None = None
@@ -296,6 +299,8 @@ class VllmConfig:
299 vllm_factors.append("None")
300 if self.structured_outputs_config:
301 vllm_factors.append(self.structured_outputs_config.compute_hash())
302 + if self.profiler_config:
303 + vllm_factors.append(self.profiler_config.compute_hash())
304 else:
305 vllm_factors.append("None")
306 vllm_factors.append(self.observability_config.compute_hash())
vllm/engine/arg_utils.py
+5 -1
@@ -50,6 +50,7 @@ from vllm.config import (
50 ObservabilityConfig,
51 ParallelConfig,
52 PoolerConfig,
53 + ProfilerConfig,
54 SchedulerConfig,
55 SpeculativeConfig,
56 StructuredOutputsConfig,
@@ -532,6 +533,8 @@ class EngineArgs:
533 worker_cls: str = ParallelConfig.worker_cls
534 worker_extension_cls: str = ParallelConfig.worker_extension_cls
535
536 + profiler_config: ProfilerConfig = get_field(VllmConfig, "profiler_config")
537 +
538 kv_transfer_config: KVTransferConfig | None = None
539 kv_events_config: KVEventsConfig | None = None
540
@@ -1164,7 +1167,7 @@ class EngineArgs:
1167 vllm_group.add_argument(
1168 "--structured-outputs-config", **vllm_kwargs["structured_outputs_config"]
1169 )
1167 -
1170 + vllm_group.add_argument("--profiler-config", **vllm_kwargs["profiler_config"])
1171 vllm_group.add_argument(
1172 "--optimization-level", **vllm_kwargs["optimization_level"]
1173 )
@@ -1782,6 +1785,7 @@ class EngineArgs:
1785 kv_transfer_config=self.kv_transfer_config,
1786 kv_events_config=self.kv_events_config,
1787 ec_transfer_config=self.ec_transfer_config,
1788 + profiler_config=self.profiler_config,
1789 additional_config=self.additional_config,
1790 optimization_level=self.optimization_level,
1791 )
vllm/entrypoints/llm.py
+17
@@ -20,6 +20,7 @@ from vllm.beam_search import (
20 from vllm.config import (
21 CompilationConfig,
22 PoolerConfig,
23 + ProfilerConfig,
24 StructuredOutputsConfig,
25 is_init_field,
26 )
@@ -211,6 +212,7 @@ class LLM:
212 structured_outputs_config: dict[str, Any]
213 | StructuredOutputsConfig
214 | None = None,
215 + profiler_config: dict[str, Any] | ProfilerConfig | None = None,
216 kv_cache_memory_bytes: int | None = None,
217 compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
218 logits_processors: list[str | type[LogitsProcessor]] | None = None,
@@ -282,6 +284,20 @@ class LLM:
284 else:
285 structured_outputs_instance = StructuredOutputsConfig()
286
287 + if profiler_config is not None:
288 + if isinstance(profiler_config, dict):
289 + profiler_config_instance = ProfilerConfig(
290 + **{
291 + k: v
292 + for k, v in profiler_config.items()
293 + if is_init_field(ProfilerConfig, k)
294 + }
295 + )
296 + else:
297 + profiler_config_instance = profiler_config
298 + else:
299 + profiler_config_instance = ProfilerConfig()
300 +
301 # warn about single-process data parallel usage.
302 _dp_size = int(kwargs.get("data_parallel_size", 1))
303 _distributed_executor_backend = kwargs.get("distributed_executor_backend")
@@ -324,6 +340,7 @@ class LLM:
340 mm_processor_kwargs=mm_processor_kwargs,
341 pooler_config=pooler_config,
342 structured_outputs_config=structured_outputs_instance,
343 + profiler_config=profiler_config_instance,
344 compilation_config=compilation_config_instance,
345 logits_processors=logits_processors,
346 **kwargs,
vllm/entrypoints/serve/profile/api_router.py
+7 -10
@@ -5,7 +5,7 @@
5 from fastapi import APIRouter, FastAPI, Request
6 from fastapi.responses import Response
7
8 -import vllm.envs as envs
8 +from vllm.config import ProfilerConfig
9 from vllm.engine.protocol import EngineClient
10 from vllm.logger import init_logger
11
@@ -35,15 +35,12 @@ async def stop_profile(raw_request: Request):
35
36
37 def attach_router(app: FastAPI):
38 - if envs.VLLM_TORCH_PROFILER_DIR:
38 + profiler_config = getattr(app.state.args, "profiler_config", None)
39 + assert profiler_config is None or isinstance(profiler_config, ProfilerConfig)
40 + if profiler_config is not None and profiler_config.profiler is not None:
41 logger.warning_once(
40 - "Torch Profiler is enabled in the API server. This should ONLY be "
41 - "used for local development!"
42 + "Profiler with mode '%s' is enabled in the "
43 + "API server. This should ONLY be used for local development!",
44 + profiler_config.profiler,
45 )
43 - elif envs.VLLM_TORCH_CUDA_PROFILE:
44 - logger.warning_once(
45 - "CUDA Profiler is enabled in the API server. This should ONLY be "
46 - "used for local development!"
47 - )
48 - if envs.VLLM_TORCH_PROFILER_DIR or envs.VLLM_TORCH_CUDA_PROFILE:
46 app.include_router(router)
vllm/envs.py
+51 -67
@@ -89,20 +89,23 @@ if TYPE_CHECKING:
89 VLLM_HTTP_TIMEOUT_KEEP_ALIVE: int = 5 # seconds
90 VLLM_PLUGINS: list[str] | None = None
91 VLLM_LORA_RESOLVER_CACHE_DIR: str | None = None
92 - VLLM_TORCH_CUDA_PROFILE: bool = False
92 + # Deprecated env variables for profiling, kept for backward compatibility
93 + # See also vllm/config/profiler.py and `--profiler-config` argument
94 + VLLM_TORCH_CUDA_PROFILE: str | None = None
95 VLLM_TORCH_PROFILER_DIR: str | None = None
94 - VLLM_TORCH_PROFILER_RECORD_SHAPES: bool = False
95 - VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
96 - VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM: bool = False
96 + VLLM_TORCH_PROFILER_RECORD_SHAPES: str | None = None
97 + VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: str | None = None
98 + VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM: str | None = None
99 + VLLM_TORCH_PROFILER_WITH_STACK: str | None = None
100 + VLLM_TORCH_PROFILER_WITH_FLOPS: str | None = None
101 + VLLM_TORCH_PROFILER_USE_GZIP: str | None = None
102 + VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL: str | None = None
103 + VLLM_PROFILER_DELAY_ITERS: str | None = None
104 + VLLM_PROFILER_MAX_ITERS: str | None = None
105 + # End of deprecated env variables for profiling
106 VLLM_USE_AOT_COMPILE: bool = False
107 VLLM_USE_BYTECODE_HOOK: bool = False
108 VLLM_FORCE_AOT_LOAD: bool = False
100 - VLLM_TORCH_PROFILER_WITH_STACK: bool = True
101 - VLLM_TORCH_PROFILER_WITH_FLOPS: bool = False
102 - VLLM_PROFILER_DELAY_ITERS: int = 0
103 - VLLM_PROFILER_MAX_ITERS: int = 0
104 - VLLM_TORCH_PROFILER_USE_GZIP: bool = True
105 - VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL: bool = True
109 VLLM_USE_TRITON_AWQ: bool = False
110 VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
111 VLLM_SKIP_P2P_CHECK: bool = False
@@ -850,71 +853,52 @@ environment_variables: dict[str, Callable[[], Any]] = {
853 "VLLM_LORA_RESOLVER_CACHE_DIR": lambda: os.getenv(
854 "VLLM_LORA_RESOLVER_CACHE_DIR", None
855 ),
853 - # Enables torch CUDA profiling if set.
854 - # On NVIDIA GPUs, this will start/stop cudaProfilerApi when triggered.
855 - "VLLM_TORCH_CUDA_PROFILE": lambda: bool(
856 - os.getenv("VLLM_TORCH_CUDA_PROFILE", "0") != "0"
857 - ),
856 + # Enables torch CUDA profiling if set to 1.
857 + # Deprecated, see profiler_config.
858 + "VLLM_TORCH_CUDA_PROFILE": lambda: os.getenv("VLLM_TORCH_CUDA_PROFILE"),
859 # Enables torch profiler if set.
859 - # Both AsyncLLM's CPU traces as well as workers'
860 - # traces (CPU & GPU) will be saved under this directory.
861 - # Note that it must be an absolute path.
862 - "VLLM_TORCH_PROFILER_DIR": lambda: (
863 - None
864 - if (val := os.getenv("VLLM_TORCH_PROFILER_DIR")) is None
865 - else (
866 - val
867 - if val.startswith("gs://") and val[5:] and val[5] != "/"
868 - else os.path.abspath(os.path.expanduser(val))
869 - )
870 - ),
871 - # Enable torch profiler to record shapes if set
872 - # VLLM_TORCH_PROFILER_RECORD_SHAPES=1. If not set, torch profiler will
873 - # not record shapes.
874 - "VLLM_TORCH_PROFILER_RECORD_SHAPES": lambda: bool(
875 - os.getenv("VLLM_TORCH_PROFILER_RECORD_SHAPES", "0") != "0"
876 - ),
877 - # Enable torch profiler to profile memory if set
878 - # VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY=1. If not set, torch profiler
879 - # will not profile memory.
880 - "VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY": lambda: bool(
881 - os.getenv("VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY", "0") != "0"
882 - ),
883 - # Enable torch profiler to profile stack if set
884 - # VLLM_TORCH_PROFILER_WITH_STACK=1. If not set, torch profiler WILL
885 - # profile stack by default.
886 - "VLLM_TORCH_PROFILER_WITH_STACK": lambda: bool(
887 - os.getenv("VLLM_TORCH_PROFILER_WITH_STACK", "1") != "0"
888 - ),
889 - # Enable torch profiler to profile flops if set
890 - # VLLM_TORCH_PROFILER_WITH_FLOPS=1. If not set, torch profiler will
891 - # not profile flops.
892 - "VLLM_TORCH_PROFILER_WITH_FLOPS": lambda: bool(
893 - os.getenv("VLLM_TORCH_PROFILER_WITH_FLOPS", "0") != "0"
894 - ),
895 - # Disable torch profiling of the AsyncLLMEngine process.
896 - # If set to 1, will not profile the engine process.
897 - "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM": lambda: bool(
898 - os.getenv("VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM", "0") != "0"
860 + # Deprecated, see profiler_config.
861 + "VLLM_TORCH_PROFILER_DIR": lambda: os.getenv("VLLM_TORCH_PROFILER_DIR"),
862 + # Enable torch profiler to record shapes if set to 1.
863 + # Deprecated, see profiler_config.
864 + "VLLM_TORCH_PROFILER_RECORD_SHAPES": lambda: (
865 + os.getenv("VLLM_TORCH_PROFILER_RECORD_SHAPES")
866 + ),
867 + # Enable torch profiler to profile memory if set to 1.
868 + # Deprecated, see profiler_config.
869 + "VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY": lambda: (
870 + os.getenv("VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY")
871 + ),
872 + # Enable torch profiler to profile stack if set to 1.
873 + # Deprecated, see profiler_config.
874 + "VLLM_TORCH_PROFILER_WITH_STACK": lambda: (
875 + os.getenv("VLLM_TORCH_PROFILER_WITH_STACK")
876 + ),
877 + # Enable torch profiler to profile flops if set to 1.
878 + # Deprecated, see profiler_config.
879 + "VLLM_TORCH_PROFILER_WITH_FLOPS": lambda: (
880 + os.getenv("VLLM_TORCH_PROFILER_WITH_FLOPS")
881 + ),
882 + # Disable torch profiling of the AsyncLLMEngine process if set to 1.
883 + # Deprecated, see profiler_config.
884 + "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM": lambda: (
885 + os.getenv("VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM")
886 ),
887 # Delay number of iterations before starting profiling when using
888 # the torch/torch CUDA profiler. If set to 0, will start profiling immediately.
902 - "VLLM_PROFILER_DELAY_ITERS": lambda: int(
903 - os.getenv("VLLM_PROFILER_DELAY_ITERS", "0")
904 - ),
889 + # Deprecated, see profiler_config.
890 + "VLLM_PROFILER_DELAY_ITERS": lambda: (os.getenv("VLLM_PROFILER_DELAY_ITERS")),
891 # Maximum number of iterations to profile when using the torch/torch CUDA profiler.
892 # If set to 0, will not limit the number of iterations.
907 - "VLLM_PROFILER_MAX_ITERS": lambda: int(os.getenv("VLLM_PROFILER_MAX_ITERS", "0")),
893 + "VLLM_PROFILER_MAX_ITERS": lambda: os.getenv("VLLM_PROFILER_MAX_ITERS"),
894 # Control whether torch profiler gzip-compresses profiling files.
909 - # Set VLLM_TORCH_PROFILER_USE_GZIP=0 to disable gzip (enabled by default).
910 - "VLLM_TORCH_PROFILER_USE_GZIP": lambda: bool(
911 - os.getenv("VLLM_TORCH_PROFILER_USE_GZIP", "1") != "0"
912 - ),
895 + # Deprecated, see profiler_config.
896 + "VLLM_TORCH_PROFILER_USE_GZIP": lambda: os.getenv("VLLM_TORCH_PROFILER_USE_GZIP"),
897 # Control whether torch profiler dumps the self_cuda_time_total table.
914 - # Set VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL=0 to disable dumping
915 - # (enabled by default).
916 - "VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL": lambda: bool(
917 - os.getenv("VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL", "1") != "0"
898 + # Set to 0 to disable dumping the table.
899 + # Deprecated, see profiler_config.
900 + "VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL": lambda: (
901 + os.getenv("VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL")
902 ),
903 # If set, vLLM will use Triton implementations of AWQ.
904 "VLLM_USE_TRITON_AWQ": lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))),
vllm/profiler/wrapper.py renamed
+47 -25
@@ -3,26 +3,27 @@
3
4 from abc import ABC, abstractmethod
5 from contextlib import nullcontext
6 +from typing import Literal
7
8 import torch
9 from typing_extensions import override
10
10 -import vllm.envs as envs
11 +from vllm.config import ProfilerConfig
12 from vllm.logger import init_logger
13
14 logger = init_logger(__name__)
15
16
17 class WorkerProfiler(ABC):
17 - def __init__(self) -> None:
18 - self._delay_iters = envs.VLLM_PROFILER_DELAY_ITERS
18 + def __init__(self, profiler_config: ProfilerConfig) -> None:
19 + self._delay_iters = profiler_config.delay_iterations
20 if self._delay_iters > 0:
21 logger.info_once(
22 "GPU profiling will start "
23 f"{self._delay_iters} steps after start_profile."
24 )
25
25 - self._max_iters = envs.VLLM_PROFILER_MAX_ITERS
26 + self._max_iters = profiler_config.max_iterations
27 if self._max_iters > 0:
28 logger.info_once(
29 "GPU profiling will stop "
@@ -133,12 +134,27 @@ class WorkerProfiler(ABC):
134 return nullcontext()
135
136
137 +TorchProfilerActivity = Literal["CPU", "CUDA", "XPU"]
138 +TorchProfilerActivityMap = {
139 + "CPU": torch.profiler.ProfilerActivity.CPU,
140 + "CUDA": torch.profiler.ProfilerActivity.CUDA,
141 + "XPU": torch.profiler.ProfilerActivity.XPU,
142 +}
143 +
144 +
145 class TorchProfilerWrapper(WorkerProfiler):
137 - def __init__(self, worker_name: str, local_rank: int) -> None:
138 - super().__init__()
146 + def __init__(
147 + self,
148 + profiler_config: ProfilerConfig,
149 + worker_name: str,
150 + local_rank: int,
151 + activities: list[TorchProfilerActivity],
152 + ) -> None:
153 + super().__init__(profiler_config)
154
155 self.local_rank = local_rank
141 - torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
156 + self.profiler_config = profiler_config
157 + torch_profiler_trace_dir = profiler_config.torch_profiler_dir
158 if local_rank in (None, 0):
159 logger.info(
160 "Torch profiling enabled. Traces will be saved to: %s",
@@ -147,24 +163,23 @@ class TorchProfilerWrapper(WorkerProfiler):
163 logger.debug(
164 "Profiler config: record_shapes=%s,"
165 "profile_memory=%s,with_stack=%s,with_flops=%s",
150 - envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
151 - envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
152 - envs.VLLM_TORCH_PROFILER_WITH_STACK,
153 - envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
166 + profiler_config.torch_profiler_record_shapes,
167 + profiler_config.torch_profiler_with_memory,
168 + profiler_config.torch_profiler_with_stack,
169 + profiler_config.torch_profiler_with_flops,
170 )
171 +
172 + self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
173 self.profiler = torch.profiler.profile(
156 - activities=[
157 - torch.profiler.ProfilerActivity.CPU,
158 - torch.profiler.ProfilerActivity.CUDA,
159 - ],
160 - record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
161 - profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
162 - with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
163 - with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
174 + activities=[TorchProfilerActivityMap[activity] for activity in activities],
175 + record_shapes=profiler_config.torch_profiler_record_shapes,
176 + profile_memory=profiler_config.torch_profiler_with_memory,
177 + with_stack=profiler_config.torch_profiler_with_stack,
178 + with_flops=profiler_config.torch_profiler_with_flops,
179 on_trace_ready=torch.profiler.tensorboard_trace_handler(
180 torch_profiler_trace_dir,
181 worker_name=worker_name,
167 - use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
182 + use_gzip=profiler_config.torch_profiler_use_gzip,
183 ),
184 )
185
@@ -176,9 +191,10 @@ class TorchProfilerWrapper(WorkerProfiler):
191 def _stop(self) -> None:
192 self.profiler.stop()
193
179 - if envs.VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL:
180 - rank = self.local_rank
181 - profiler_dir = envs.VLLM_TORCH_PROFILER_DIR
194 + profiler_config = self.profiler_config
195 + rank = self.local_rank
196 + if profiler_config.torch_profiler_dump_cuda_time_total:
197 + profiler_dir = profiler_config.torch_profiler_dir
198 profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
199 sort_key = "self_cuda_time_total"
200 table = self.profiler.key_averages().table(sort_by=sort_key)
@@ -189,6 +205,12 @@ class TorchProfilerWrapper(WorkerProfiler):
205 # only print profiler results on rank 0
206 if rank == 0:
207 print(table)
208 + if self.dump_cpu_time_total and rank == 0:
209 + logger.info(
210 + self.profiler.key_averages().table(
211 + sort_by="self_cpu_time_total", row_limit=50
212 + )
213 + )
214
215 @override
216 def annotate_context_manager(self, name: str):
@@ -196,8 +218,8 @@ class TorchProfilerWrapper(WorkerProfiler):
218
219
220 class CudaProfilerWrapper(WorkerProfiler):
199 - def __init__(self) -> None:
200 - super().__init__()
221 + def __init__(self, profiler_config: ProfilerConfig) -> None:
222 + super().__init__(profiler_config)
223 # Note: lazy import to avoid dependency issues if CUDA is not available.
224 import torch.cuda.profiler as cuda_profiler
225
vllm/v1/engine/async_llm.py
+7 -15
@@ -166,32 +166,24 @@ class AsyncLLM(EngineClient):
166 pass
167
168 if (
169 - envs.VLLM_TORCH_PROFILER_DIR
170 - and not envs.VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM
169 + vllm_config.profiler_config.profiler == "torch"
170 + and not vllm_config.profiler_config.ignore_frontend
171 ):
172 + profiler_dir = vllm_config.profiler_config.torch_profiler_dir
173 logger.info(
174 "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501
174 - envs.VLLM_TORCH_PROFILER_DIR,
175 + profiler_dir,
176 )
176 - if envs.VLLM_PROFILER_MAX_ITERS > 0 or envs.VLLM_PROFILER_DELAY_ITERS > 0:
177 - logger.warning_once(
178 - "Torch profiler received max_iters or delay_iters setting. These "
179 - "are not compatible with the AsyncLLM profiler and will be ignored "
180 - "for the AsyncLLM process. Engine process profiling will still "
181 - "respect these settings. Consider setting "
182 - "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM=1 to disable "
183 - "AsyncLLM profiling."
184 - )
177 worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
178 self.profiler = torch.profiler.profile(
179 activities=[
180 torch.profiler.ProfilerActivity.CPU,
181 ],
190 - with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
182 + with_stack=vllm_config.profiler_config.torch_profiler_with_stack,
183 on_trace_ready=torch.profiler.tensorboard_trace_handler(
192 - envs.VLLM_TORCH_PROFILER_DIR,
184 + profiler_dir,
185 worker_name=worker_name,
194 - use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
186 + use_gzip=vllm_config.profiler_config.torch_profiler_use_gzip,
187 ),
188 )
189 else:
vllm/v1/worker/cpu_worker.py
+9 -27
@@ -13,6 +13,7 @@ from vllm.logger import init_logger
13 from vllm.model_executor.utils import set_random_seed
14 from vllm.platforms import CpuArchEnum, current_platform
15 from vllm.platforms.cpu import CpuPlatform, LogicalCPUInfo
16 +from vllm.profiler.wrapper import TorchProfilerWrapper
17 from vllm.v1.worker.cpu_model_runner import CPUModelRunner
18 from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
19
@@ -38,30 +39,17 @@ class CPUWorker(Worker):
39
40 self.parallel_config.disable_custom_all_reduce = True
41
41 - # Torch profiler. Enabled and configured through env vars:
42 - # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
42 + # Torch profiler. Enabled and configured through profiler_config.
43 self.profiler: Any | None = None
44 - if envs.VLLM_TORCH_PROFILER_DIR:
45 - torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
44 + profiler_config = vllm_config.profiler_config
45 + if profiler_config.profiler == "torch":
46 worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
47 - logger.info(
48 - "Profiling enabled. Traces will be saved to: %s",
49 - torch_profiler_trace_dir,
47 + self.profiler = TorchProfilerWrapper(
48 + profiler_config,
49 + worker_name=worker_name,
50 + local_rank=self.local_rank,
51 + activities=["CPU"],
52 )
51 - self.profiler = torch.profiler.profile(
52 - activities=[
53 - torch.profiler.ProfilerActivity.CPU,
54 - ],
55 - record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
56 - profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
57 - with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
58 - with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
59 - on_trace_ready=torch.profiler.tensorboard_trace_handler(
60 - torch_profiler_trace_dir, worker_name=worker_name, use_gzip=False
61 - ),
62 - )
63 - else:
64 - self.profiler = None
53
54 def init_device(self):
55 # Setup OpenMP threads affinity.
@@ -202,9 +190,3 @@ class CPUWorker(Worker):
190 self.profiler.start()
191 else:
192 self.profiler.stop()
205 - if self.local_rank == 0:
206 - logger.info(
207 - self.profiler.key_averages().table(
208 - sort_by="self_cpu_time_total", row_limit=50
209 - )
210 - )
vllm/v1/worker/gpu_worker.py
+10 -8
@@ -38,7 +38,7 @@ from vllm.model_executor import set_random_seed
38 from vllm.model_executor.models.interfaces import is_mixture_of_experts
39 from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
40 from vllm.platforms import current_platform
41 -from vllm.profiler.gpu_profiler import CudaProfilerWrapper, TorchProfilerWrapper
41 +from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper
42 from vllm.sequence import IntermediateTensors
43 from vllm.tasks import SupportedTask
44 from vllm.utils.mem_constants import GiB_bytes
@@ -92,17 +92,19 @@ class Worker(WorkerBase):
92 # Buffers saved before sleep
93 self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
94
95 - # Torch/CUDA profiler. Enabled and configured through env vars:
96 - # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
97 - # VLLM_TORCH_CUDA_PROFILE=1
95 + # Torch/CUDA profiler. Enabled and configured through profiler_config.
96 self.profiler: Any | None = None
99 - if envs.VLLM_TORCH_PROFILER_DIR:
97 + profiler_config = vllm_config.profiler_config
98 + if profiler_config.profiler == "torch":
99 worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
100 self.profiler = TorchProfilerWrapper(
102 - worker_name=worker_name, local_rank=self.local_rank
101 + profiler_config,
102 + worker_name=worker_name,
103 + local_rank=self.local_rank,
104 + activities=["CPU", "CUDA"],
105 )
104 - elif envs.VLLM_TORCH_CUDA_PROFILE:
105 - self.profiler = CudaProfilerWrapper()
106 + elif profiler_config.profiler == "cuda":
107 + self.profiler = CudaProfilerWrapper(profiler_config)
108 else:
109 self.profiler = None
110
vllm/v1/worker/tpu_worker.py
+2 -2
@@ -98,10 +98,10 @@ class TPUWorker:
98 # MP runtime is initialized.
99 self.profiler = None
100 self.profile_dir = None
101 - if envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1:
101 + if vllm_config.profiler_config.profiler == "torch" and self.rank < 1:
102 # For TPU, we can only have 1 active profiler session for 1 profiler
103 # server. So we only profile on rank0.
104 - self.profile_dir = envs.VLLM_TORCH_PROFILER_DIR
104 + self.profile_dir = vllm_config.profiler_config.torch_profiler_dir
105 logger.info(
106 "Profiling enabled. Traces will be saved to: %s", self.profile_dir
107 )
vllm/v1/worker/xpu_worker.py
+9 -33
@@ -6,12 +6,12 @@ from typing import Any
6 import torch
7 import torch.distributed
8
9 -import vllm.envs as envs
9 from vllm.config import VllmConfig
10 from vllm.distributed import get_world_group
11 from vllm.logger import init_logger
12 from vllm.model_executor import set_random_seed
13 from vllm.platforms import current_platform
14 +from vllm.profiler.wrapper import TorchProfilerWrapper
15 from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
16 from vllm.v1.worker.xpu_model_runner import XPUModelRunner
17
@@ -36,41 +36,17 @@ class XPUWorker(Worker):
36 assert device_config.device_type == "xpu"
37 assert current_platform.is_xpu()
38
39 - # Torch profiler. Enabled and configured through env vars:
40 - # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
39 + # Torch profiler. Enabled and configured through profiler_config.
40 self.profiler: Any | None = None
42 - if envs.VLLM_TORCH_PROFILER_DIR:
43 - torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
41 + profiler_config = vllm_config.profiler_config
42 + if profiler_config.profiler == "torch":
43 worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
45 - logger.info(
46 - "Profiling enabled. Traces will be saved to: %s",
47 - torch_profiler_trace_dir,
44 + self.profiler = TorchProfilerWrapper(
45 + profiler_config,
46 + worker_name=worker_name,
47 + local_rank=self.local_rank,
48 + activities=["CPU", "XPU"],
49 )
49 - logger.debug(
50 - "Profiler config: record_shapes=%s,"
51 - "profile_memory=%s,with_stack=%s,with_flops=%s",
52 - envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
53 - envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
54 - envs.VLLM_TORCH_PROFILER_WITH_STACK,
55 - envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
56 - )
57 - self.profiler = torch.profiler.profile(
58 - activities=[
59 - torch.profiler.ProfilerActivity.CPU,
60 - torch.profiler.ProfilerActivity.XPU,
61 - ],
62 - record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
63 - profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
64 - with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
65 - with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
66 - on_trace_ready=torch.profiler.tensorboard_trace_handler(
67 - torch_profiler_trace_dir,
68 - worker_name=worker_name,
69 - use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
70 - ),
71 - )
72 - else:
73 - self.profiler = None
50
51 # we provide this function due to `torch.xpu.mem_get_info()` doesn't
52 # return correct free_gpu_memory on intel client GPU. We need to