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, 23:30 UTC
[Bugfix][Platform][CPU] Fix cuda platform detection on CPU backend edge case (#13358)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Feb 17, 2025 at 02:55 UTC
d67cc21b787dc594f6b168c4e44a0c6ecb415385
1 file changed
+9
-2
vllm/platforms/__init__.py
+9
-2
@@ -33,12 +33,19 @@ def cuda_platform_plugin() -> Optional[str]:
33
is_cuda = False
34
35
try:
36
+ from importlib.metadata import version
37
+
38
from vllm.utils import import_pynvml
39
pynvml = import_pynvml()
40
pynvml.nvmlInit()
41
try:
40
- if pynvml.nvmlDeviceGetCount() > 0:
41
- is_cuda = True
42
+ # NOTE: Edge case: vllm cpu build on a GPU machine.
43
+ # Third-party pynvml can be imported in cpu build,
44
+ # we need to check if vllm is built with cpu too.
45
+ # Otherwise, vllm will always activate cuda plugin
46
+ # on a GPU machine, even if in a cpu build.
47
+ is_cuda = (pynvml.nvmlDeviceGetCount() > 0
48
+ and "cpu" not in version("vllm"))
49
finally:
50
pynvml.nvmlShutdown()
51
except Exception as e: