Skip to content

Commit b37cdce

Browse files
authored
[Core] Cache some utils (#3474)
1 parent b30880a commit b37cdce

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

vllm/utils.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import subprocess
55
import uuid
66
import gc
7+
from functools import cache
78
from platform import uname
89
from typing import List, Tuple, Union
910
from packaging.version import parse, Version
@@ -120,6 +121,7 @@ def is_hip() -> bool:
120121
return torch.version.hip is not None
121122

122123

124+
@cache
123125
def is_neuron() -> bool:
124126
try:
125127
import transformers_neuronx
@@ -128,6 +130,7 @@ def is_neuron() -> bool:
128130
return transformers_neuronx is not None
129131

130132

133+
@cache
131134
def get_max_shared_memory_bytes(gpu: int = 0) -> int:
132135
"""Returns the maximum shared memory per thread block in bytes."""
133136
# NOTE: This import statement should be executed lazily since
@@ -151,6 +154,7 @@ def random_uuid() -> str:
151154
return str(uuid.uuid4().hex)
152155

153156

157+
@cache
154158
def in_wsl() -> bool:
155159
# Reference: https://github.com/microsoft/WSL/issues/4071
156160
return "microsoft" in " ".join(uname()).lower()
@@ -225,6 +229,7 @@ def set_cuda_visible_devices(device_ids: List[int]) -> None:
225229
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, device_ids))
226230

227231

232+
@cache
228233
def get_nvcc_cuda_version() -> Optional[Version]:
229234
cuda_home = os.environ.get('CUDA_HOME')
230235
if not cuda_home:

0 commit comments

Comments
 (0)