mirror of
https://github.com/wassname/vllm.git
synced 2026-06-27 18:45:36 +08:00
TJian
6ccc0bfffb
Co-authored-by: Philipp Moritz <pcmoritz@gmail.com> Co-authored-by: Amir Balwel <amoooori04@gmail.com> Co-authored-by: root <kuanfu.liu@akirakan.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: kuanfu <kuanfu.liu@embeddedllm.com> Co-authored-by: miloice <17350011+kliuae@users.noreply.github.com>
60 lines
1.4 KiB
Python
60 lines
1.4 KiB
Python
import enum
|
|
import uuid
|
|
from platform import uname
|
|
|
|
import psutil
|
|
import torch
|
|
|
|
from vllm._C import cuda_utils
|
|
|
|
|
|
class Device(enum.Enum):
|
|
GPU = enum.auto()
|
|
CPU = enum.auto()
|
|
|
|
|
|
class Counter:
|
|
|
|
def __init__(self, start: int = 0) -> None:
|
|
self.counter = start
|
|
|
|
def __next__(self) -> int:
|
|
i = self.counter
|
|
self.counter += 1
|
|
return i
|
|
|
|
def reset(self) -> None:
|
|
self.counter = 0
|
|
|
|
|
|
def is_hip() -> bool:
|
|
return torch.version.hip is not None
|
|
|
|
|
|
def get_max_shared_memory_bytes(gpu: int = 0) -> int:
|
|
"""Returns the maximum shared memory per thread block in bytes."""
|
|
# https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html
|
|
cudaDevAttrMaxSharedMemoryPerBlockOptin = 97 if not is_hip() else 74
|
|
max_shared_mem = cuda_utils.get_device_attribute(
|
|
cudaDevAttrMaxSharedMemoryPerBlockOptin, gpu)
|
|
return int(max_shared_mem)
|
|
|
|
|
|
def get_gpu_memory(gpu: int = 0) -> int:
|
|
"""Returns the total memory of the GPU in bytes."""
|
|
return torch.cuda.get_device_properties(gpu).total_memory
|
|
|
|
|
|
def get_cpu_memory() -> int:
|
|
"""Returns the total CPU memory of the node in bytes."""
|
|
return psutil.virtual_memory().total
|
|
|
|
|
|
def random_uuid() -> str:
|
|
return str(uuid.uuid4().hex)
|
|
|
|
|
|
def in_wsl() -> bool:
|
|
# Reference: https://github.com/microsoft/WSL/issues/4071
|
|
return "microsoft" in " ".join(uname()).lower()
|