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 15, 16:00 UTC
[misc] instruct pytorch to use nvml-based cuda check (#15951)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Apr 3, 2025 at 01:02 UTC
1cab43c2d2172459987652a7fbf922c6611a71fa
2 files changed
+25
-16
vllm/__init__.py
+4
-16
@@ -4,9 +4,10 @@
4
# version library first. Such assumption is critical for some customization.
5
from .version import __version__, __version_tuple__ # isort:skip
6
7
-import os
8
-
9
-import torch
7
+# The environment variables override should be imported before any other
8
+# modules to ensure that the environment variables are set before any
9
+# other modules are imported.
10
+import vllm.env_override # isort:skip # noqa: F401
11
12
from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
13
from vllm.engine.async_llm_engine import AsyncLLMEngine
@@ -23,19 +24,6 @@ from vllm.outputs import (ClassificationOutput, ClassificationRequestOutput,
24
from vllm.pooling_params import PoolingParams
25
from vllm.sampling_params import SamplingParams
26
26
-# set some common config/environment variables that should be set
27
-# for all processes created by vllm and all processes
28
-# that interact with vllm workers.
29
-# they are executed whenever `import vllm` is called.
30
-
31
-# see https://github.com/NVIDIA/nccl/issues/1234
32
-os.environ['NCCL_CUMEM_ENABLE'] = '0'
33
-
34
-# see https://github.com/vllm-project/vllm/issues/10480
35
-os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = '1'
36
-# see https://github.com/vllm-project/vllm/issues/10619
37
-torch._inductor.config.compile_threads = 1
38
-
27
__all__ = [
28
"__version__",
29
"__version_tuple__",
vllm/env_override.py
new
+21
@@ -0,0 +1,21 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import os
3
+
4
+import torch
5
+
6
+# set some common config/environment variables that should be set
7
+# for all processes created by vllm and all processes
8
+# that interact with vllm workers.
9
+# they are executed whenever `import vllm` is called.
10
+
11
+# see https://github.com/NVIDIA/nccl/issues/1234
12
+os.environ['NCCL_CUMEM_ENABLE'] = '0'
13
+
14
+# see https://github.com/vllm-project/vllm/pull/15951
15
+# it avoids unintentional cuda initialization from torch.cuda.is_available()
16
+os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
17
+
18
+# see https://github.com/vllm-project/vllm/issues/10480
19
+os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = '1'
20
+# see https://github.com/vllm-project/vllm/issues/10619
21
+torch._inductor.config.compile_threads = 1