Skip to content

Commit

Permalink
[Bugfix][Platform][CPU] Fix cuda platform detection on CPU backend ed…
Browse files Browse the repository at this point in the history
…ge case (#13358)

Signed-off-by: Isotr0py <2037008807@qq.com>
  • Loading branch information
Isotr0py authored Feb 16, 2025
1 parent e18227b commit d67cc21
Showing 1 changed file with 9 additions and 2 deletions.
11 changes: 9 additions & 2 deletions vllm/platforms/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,19 @@ def cuda_platform_plugin() -> Optional[str]:
is_cuda = False

try:
from importlib.metadata import version

from vllm.utils import import_pynvml
pynvml = import_pynvml()
pynvml.nvmlInit()
try:
if pynvml.nvmlDeviceGetCount() > 0:
is_cuda = True
# NOTE: Edge case: vllm cpu build on a GPU machine.
# Third-party pynvml can be imported in cpu build,
# we need to check if vllm is built with cpu too.
# Otherwise, vllm will always activate cuda plugin
# on a GPU machine, even if in a cpu build.
is_cuda = (pynvml.nvmlDeviceGetCount() > 0
and "cpu" not in version("vllm"))
finally:
pynvml.nvmlShutdown()
except Exception as e:
Expand Down

0 comments on commit d67cc21

Please sign in to comment.