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
Add `VLLM_NVTX_SCOPES_FOR_PROFILING=1` to enable `nvtx.annotate` scopes (#25501)
Signed-off-by: Corey Lowman <clowman1993@gmail.com>
Corey Lowman committed
Sep 23, 2025 at 18:50 UTC
c85d75cf08330f16f6c076d9cb1a586fa6c6fbb8
2 files changed
+22
-3
vllm/envs.py
+5
@@ -187,6 +187,7 @@ if TYPE_CHECKING:
187
VLLM_DISABLE_PAD_FOR_CUDAGRAPH: bool = False
188
VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False
189
VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False
190
+ VLLM_NVTX_SCOPES_FOR_PROFILING: bool = False
191
VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True
192
VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "VLLM_OBJECT_STORAGE_SHM_BUFFER"
193
VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024
@@ -1387,6 +1388,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1388
"VLLM_CUSTOM_SCOPES_FOR_PROFILING":
1389
lambda: bool(int(os.getenv("VLLM_CUSTOM_SCOPES_FOR_PROFILING", "0"))),
1390
1391
+ # Add optional nvtx scopes for profiling, disable to avoid overheads
1392
+ "VLLM_NVTX_SCOPES_FOR_PROFILING":
1393
+ lambda: bool(int(os.getenv("VLLM_NVTX_SCOPES_FOR_PROFILING", "0"))),
1394
+
1395
# Represent block hashes in KV cache events as 64-bit integers instead of
1396
# raw bytes. Defaults to True for backward compatibility.
1397
"VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES":
vllm/v1/utils.py
+17
-3
@@ -375,8 +375,22 @@ def report_usage_stats(
375
})
376
377
378
+_PROFILER_FUNC = None
379
+
380
+
381
def record_function_or_nullcontext(name: str) -> AbstractContextManager:
382
+ global _PROFILER_FUNC
383
+
384
+ # fast path assume it is set
385
+ if _PROFILER_FUNC is not None:
386
+ return _PROFILER_FUNC(name)
387
+
388
+ func = contextlib.nullcontext
389
if envs.VLLM_CUSTOM_SCOPES_FOR_PROFILING:
380
- return record_function(name)
381
- else:
382
- return contextlib.nullcontext()
390
+ func = record_function
391
+ elif envs.VLLM_NVTX_SCOPES_FOR_PROFILING:
392
+ import nvtx
393
+ func = nvtx.annotate
394
+
395
+ _PROFILER_FUNC = func
396
+ return func(name)