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 14, 05:30 UTC
[Core] Add torch profiler CPU traces for AsyncLLM. (#21794)
Signed-off-by: Chenheli Hua <huachenheli@outlook.com>
Chenheli Hua committed
Aug 19, 2025 at 19:32 UTC
e58c5a97688750e7930f13b6fe556d9a28a5b2d9
2 files changed
+35
-4
vllm/envs.py
+4
-2
@@ -667,8 +667,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
667
"VLLM_LORA_RESOLVER_CACHE_DIR":
668
lambda: os.getenv("VLLM_LORA_RESOLVER_CACHE_DIR", None),
669
670
- # Enables torch profiler if set. Path to the directory where torch profiler
671
- # traces are saved. Note that it must be an absolute path.
670
+ # Enables torch profiler if set.
671
+ # Both AsyncLLM's CPU traces as well as workers'
672
+ # traces (CPU & GPU) will be saved under this directory.
673
+ # Note that it must be an absolute path.
674
"VLLM_TORCH_PROFILER_DIR":
675
lambda: (None if os.getenv("VLLM_TORCH_PROFILER_DIR", None) is None else os
676
.path.expanduser(os.getenv("VLLM_TORCH_PROFILER_DIR", "."))),
vllm/v1/engine/async_llm.py
+31
-2
@@ -1,12 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import asyncio
4
+import os
5
+import socket
6
import time
7
from collections.abc import AsyncGenerator, Iterable, Mapping
8
from copy import copy
9
from typing import Any, Optional, Union
10
11
import numpy as np
12
+import torch
13
14
import vllm.envs as envs
15
from vllm.config import ModelConfig, VllmConfig
@@ -144,6 +147,26 @@ class AsyncLLM(EngineClient):
147
except RuntimeError:
148
pass
149
150
+ if envs.VLLM_TORCH_PROFILER_DIR:
151
+ logger.info(
152
+ "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501
153
+ envs.VLLM_TORCH_PROFILER_DIR)
154
+ worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
155
+ self.profiler = torch.profiler.profile(
156
+ activities=[
157
+ torch.profiler.ProfilerActivity.CPU,
158
+ ],
159
+ with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
160
+ on_trace_ready=torch.profiler.tensorboard_trace_handler(
161
+ envs.VLLM_TORCH_PROFILER_DIR,
162
+ worker_name=worker_name,
163
+ use_gzip=True))
164
+ else:
165
+ logger.info(
166
+ "Torch profiler disabled. AsyncLLM CPU traces will not be collected." # noqa: E501
167
+ )
168
+ self.profiler = None
169
+
170
@classmethod
171
@deprecate_kwargs(
172
"disable_log_requests",
@@ -562,10 +585,16 @@ class AsyncLLM(EngineClient):
585
raise self.dead_error
586
587
async def start_profile(self) -> None:
565
- await self.engine_core.profile_async(True)
588
+ coros = [self.engine_core.profile_async(True)]
589
+ if self.profiler is not None:
590
+ coros.append(asyncio.to_thread(self.profiler.start))
591
+ await asyncio.gather(*coros)
592
593
async def stop_profile(self) -> None:
568
- await self.engine_core.profile_async(False)
594
+ coros = [self.engine_core.profile_async(False)]
595
+ if self.profiler is not None:
596
+ coros.append(asyncio.to_thread(self.profiler.stop))
597
+ await asyncio.gather(*coros)
598
599
async def reset_mm_cache(self) -> None:
600
self.processor.mm_registry.reset_processor_cache(self.model_config)