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, 22:35 UTC
Expose PyTorch profiler configuration to environment variables (#21803)
Signed-off-by: Csrayz <33659823+Csrayz@users.noreply.github.com>
Csrayz committed
Jul 30, 2025 at 10:46 UTC
b917da442b820245f537602d752e7146e66dd37a
4 files changed
+60
-4
docs/contributing/profiling.md
+6
-1
@@ -5,7 +5,12 @@
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/`
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:
9
+
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
15
The OpenAI server also needs to be started with the `VLLM_TORCH_PROFILER_DIR` environment variable set.
16
vllm/envs.py
+29
@@ -80,6 +80,10 @@ if TYPE_CHECKING:
80
VLLM_PLUGINS: Optional[list[str]] = None
81
VLLM_LORA_RESOLVER_CACHE_DIR: Optional[str] = None
82
VLLM_TORCH_PROFILER_DIR: Optional[str] = None
83
+ VLLM_TORCH_PROFILER_RECORD_SHAPES: bool = False
84
+ VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
85
+ VLLM_TORCH_PROFILER_WITH_STACK: bool = True
86
+ VLLM_TORCH_PROFILER_WITH_FLOPS: bool = False
87
VLLM_USE_TRITON_AWQ: bool = False
88
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
89
VLLM_SKIP_P2P_CHECK: bool = False
@@ -629,6 +633,31 @@ environment_variables: dict[str, Callable[[], Any]] = {
633
lambda: (None if os.getenv("VLLM_TORCH_PROFILER_DIR", None) is None else os
634
.path.expanduser(os.getenv("VLLM_TORCH_PROFILER_DIR", "."))),
635
636
+ # Enable torch profiler to record shapes if set
637
+ # VLLM_TORCH_PROFILER_RECORD_SHAPES=1. If not set, torch profiler will
638
+ # not record shapes.
639
+ "VLLM_TORCH_PROFILER_RECORD_SHAPES":
640
+ lambda: bool(os.getenv("VLLM_TORCH_PROFILER_RECORD_SHAPES", "0") != "0"),
641
+
642
+ # Enable torch profiler to profile memory if set
643
+ # VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY=1. If not set, torch profiler
644
+ # will not profile memory.
645
+ "VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY":
646
+ lambda: bool(
647
+ os.getenv("VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY", "0") != "0"),
648
+
649
+ # Enable torch profiler to profile stack if set
650
+ # VLLM_TORCH_PROFILER_WITH_STACK=1. If not set, torch profiler WILL
651
+ # profile stack by default.
652
+ "VLLM_TORCH_PROFILER_WITH_STACK":
653
+ lambda: bool(os.getenv("VLLM_TORCH_PROFILER_WITH_STACK", "1") != "0"),
654
+
655
+ # Enable torch profiler to profile flops if set
656
+ # VLLM_TORCH_PROFILER_WITH_FLOPS=1. If not set, torch profiler will
657
+ # not profile flops.
658
+ "VLLM_TORCH_PROFILER_WITH_FLOPS":
659
+ lambda: bool(os.getenv("VLLM_TORCH_PROFILER_WITH_FLOPS", "0") != "0"),
660
+
661
# If set, vLLM will use Triton implementations of AWQ.
662
"VLLM_USE_TRITON_AWQ":
663
lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))),
vllm/v1/worker/gpu_worker.py
+13
-2
@@ -71,12 +71,23 @@ class Worker(WorkerBase):
71
torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
72
logger.info("Profiling enabled. Traces will be saved to: %s",
73
torch_profiler_trace_dir)
74
+ logger.debug(
75
+ "Profiler config: record_shapes=%s,"
76
+ "profile_memory=%s,with_stack=%s,with_flops=%s",
77
+ envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
78
+ envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
79
+ envs.VLLM_TORCH_PROFILER_WITH_STACK,
80
+ envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
81
+ )
82
self.profiler = torch.profiler.profile(
83
activities=[
84
torch.profiler.ProfilerActivity.CPU,
85
torch.profiler.ProfilerActivity.CUDA,
86
],
79
- with_stack=True,
87
+ record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
88
+ profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
89
+ with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
90
+ with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
91
on_trace_ready=torch.profiler.tensorboard_trace_handler(
92
torch_profiler_trace_dir, use_gzip=True))
93
else:
@@ -209,7 +220,7 @@ class Worker(WorkerBase):
220
221
@torch.inference_mode()
222
def determine_available_memory(self) -> int:
212
- """Profiles the peak memory usage of the model to determine how much
223
+ """Profiles the peak memory usage of the model to determine how much
224
memory can be used for KV cache without OOMs.
225
226
The engine will first conduct a profiling of the existing memory usage.
vllm/v1/worker/xpu_worker.py
+12
-1
@@ -41,12 +41,23 @@ class XPUWorker(Worker):
41
torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
42
logger.info("Profiling enabled. Traces will be saved to: %s",
43
torch_profiler_trace_dir)
44
+ logger.debug(
45
+ "Profiler config: record_shapes=%s,"
46
+ "profile_memory=%s,with_stack=%s,with_flops=%s",
47
+ envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
48
+ envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
49
+ envs.VLLM_TORCH_PROFILER_WITH_STACK,
50
+ envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
51
+ )
52
self.profiler = torch.profiler.profile(
53
activities=[
54
torch.profiler.ProfilerActivity.CPU,
55
torch.profiler.ProfilerActivity.XPU,
56
],
49
- with_stack=True,
57
+ record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
58
+ profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
59
+ with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
60
+ with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
61
on_trace_ready=torch.profiler.tensorboard_trace_handler(
62
torch_profiler_trace_dir, use_gzip=True))
63
else: