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 12, 07:15 UTC
Make PyTorch profiler gzip and CUDA time dump configurable (#29568)
Signed-off-by: Yifei Zhang <yifei.zhang1992@outlook.com>
Yifei Zhang committed
Dec 1, 2025 at 12:30 UTC
1ab8fc8197f5aa0ae471ef0ff1c725b733594f7a
5 files changed
+35
-13
docs/contributing/profiling.md
+2
@@ -11,6 +11,8 @@ We support tracing vLLM workers using the `torch.profiler` module. You can enabl
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.
18
vllm/envs.py
+13
@@ -100,6 +100,8 @@ if TYPE_CHECKING:
100
VLLM_TORCH_PROFILER_WITH_FLOPS: bool = False
101
VLLM_PROFILER_DELAY_ITERS: int = 0
102
VLLM_PROFILER_MAX_ITERS: int = 0
103
+ VLLM_TORCH_PROFILER_USE_GZIP: bool = True
104
+ VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL: bool = True
105
VLLM_USE_TRITON_AWQ: bool = False
106
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
107
VLLM_SKIP_P2P_CHECK: bool = False
@@ -890,6 +892,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
892
# Maximum number of iterations to profile when using the torch/torch CUDA profiler.
893
# If set to 0, will not limit the number of iterations.
894
"VLLM_PROFILER_MAX_ITERS": lambda: int(os.getenv("VLLM_PROFILER_MAX_ITERS", "0")),
895
+ # Control whether torch profiler gzip-compresses profiling files.
896
+ # Set VLLM_TORCH_PROFILER_USE_GZIP=0 to disable gzip (enabled by default).
897
+ "VLLM_TORCH_PROFILER_USE_GZIP": lambda: bool(
898
+ os.getenv("VLLM_TORCH_PROFILER_USE_GZIP", "1") != "0"
899
+ ),
900
+ # Control whether torch profiler dumps the self_cuda_time_total table.
901
+ # Set VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL=0 to disable dumping
902
+ # (enabled by default).
903
+ "VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL": lambda: bool(
904
+ os.getenv("VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL", "1") != "0"
905
+ ),
906
# If set, vLLM will use Triton implementations of AWQ.
907
"VLLM_USE_TRITON_AWQ": lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))),
908
# If set, allow loading or unloading lora adapters in runtime,
vllm/profiler/gpu_profiler.py
+14
-11
@@ -162,7 +162,9 @@ class TorchProfilerWrapper(WorkerProfiler):
162
with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
163
with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
164
on_trace_ready=torch.profiler.tensorboard_trace_handler(
165
- torch_profiler_trace_dir, worker_name=worker_name, use_gzip=True
165
+ torch_profiler_trace_dir,
166
+ worker_name=worker_name,
167
+ use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
168
),
169
)
170
@@ -174,18 +176,19 @@ class TorchProfilerWrapper(WorkerProfiler):
176
def _stop(self) -> None:
177
self.profiler.stop()
178
177
- rank = self.local_rank
178
- profiler_dir = envs.VLLM_TORCH_PROFILER_DIR
179
- profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
180
- sort_key = "self_cuda_time_total"
181
- table = self.profiler.key_averages().table(sort_by=sort_key)
179
+ if envs.VLLM_TORCH_PROFILER_DUMP_CUDA_TIME_TOTAL:
180
+ rank = self.local_rank
181
+ profiler_dir = envs.VLLM_TORCH_PROFILER_DIR
182
+ profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
183
+ sort_key = "self_cuda_time_total"
184
+ table = self.profiler.key_averages().table(sort_by=sort_key)
185
183
- with open(profiler_out_file, "w") as f:
184
- print(table, file=f)
186
+ with open(profiler_out_file, "w") as f:
187
+ print(table, file=f)
188
186
- # only print profiler results on rank 0
187
- if rank == 0:
188
- print(table)
189
+ # only print profiler results on rank 0
190
+ if rank == 0:
191
+ print(table)
192
193
@override
194
def annotate_context_manager(self, name: str):
vllm/v1/engine/async_llm.py
+3
-1
@@ -190,7 +190,9 @@ class AsyncLLM(EngineClient):
190
],
191
with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
192
on_trace_ready=torch.profiler.tensorboard_trace_handler(
193
- envs.VLLM_TORCH_PROFILER_DIR, worker_name=worker_name, use_gzip=True
193
+ envs.VLLM_TORCH_PROFILER_DIR,
194
+ worker_name=worker_name,
195
+ use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
196
),
197
)
198
else:
vllm/v1/worker/xpu_worker.py
+3
-1
@@ -64,7 +64,9 @@ class XPUWorker(Worker):
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, worker_name=worker_name, use_gzip=True
67
+ torch_profiler_trace_dir,
68
+ worker_name=worker_name,
69
+ use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
70
),
71
)
72
else: