[MPS] Support sglang.check_env (#20753)
Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
This commit is contained in:
@@ -10,7 +10,7 @@ from collections import OrderedDict, defaultdict
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils import is_hip, is_musa, is_npu
|
||||
from sglang.srt.utils import is_hip, is_mps, is_musa, is_npu
|
||||
|
||||
|
||||
def is_cuda_v2():
|
||||
@@ -512,6 +512,82 @@ class MUSAEnv(BaseEnv):
|
||||
return {}
|
||||
|
||||
|
||||
class MPSEnv(BaseEnv):
|
||||
"""Environment checker for Apple Silicon MPS"""
|
||||
|
||||
EXTRA_PACKAGE_LIST = ["mlx", "mlx-lm", "mlx-metal"]
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.package_list.extend(MPSEnv.EXTRA_PACKAGE_LIST)
|
||||
|
||||
def get_info(self):
|
||||
import platform
|
||||
|
||||
info = {"MPS available": torch.backends.mps.is_available()}
|
||||
if not info["MPS available"]:
|
||||
return info
|
||||
|
||||
info["macOS Version"] = platform.mac_ver()[0]
|
||||
|
||||
try:
|
||||
info["macOS Build"] = subprocess.check_output(
|
||||
["sw_vers", "-buildVersion"], text=True
|
||||
).strip()
|
||||
except Exception:
|
||||
info["macOS Build"] = "Not Available"
|
||||
|
||||
for label, key in [
|
||||
("Apple Silicon", "machdep.cpu.brand_string"),
|
||||
("Unified Memory", "hw.memsize"),
|
||||
("CPU Cores (Total)", "hw.ncpu"),
|
||||
]:
|
||||
try:
|
||||
info[label] = subprocess.check_output(
|
||||
["sysctl", "-n", key], text=True
|
||||
).strip()
|
||||
except Exception:
|
||||
info[label] = "Not Available"
|
||||
|
||||
try:
|
||||
mem_bytes = int(info["Unified Memory"])
|
||||
info["Unified Memory"] = f"{mem_bytes / 1024**3:.1f} GB"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for label, key in [
|
||||
("CPU Cores (Performance)", "hw.perflevel0.logicalcpu"),
|
||||
("CPU Cores (Efficiency)", "hw.perflevel1.logicalcpu"),
|
||||
]:
|
||||
try:
|
||||
info[label] = subprocess.check_output(
|
||||
["sysctl", "-n", key], text=True
|
||||
).strip()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Single system_profiler call for both Metal support and GPU cores
|
||||
info["Metal Support"] = "Not Available"
|
||||
info["GPU Cores"] = "Not Available"
|
||||
try:
|
||||
sp = subprocess.check_output(
|
||||
["system_profiler", "SPDisplaysDataType"], text=True
|
||||
)
|
||||
for line in sp.splitlines():
|
||||
line = line.strip()
|
||||
if "Metal Support" in line or "Metal Family" in line:
|
||||
info["Metal Support"] = line.partition(":")[2].strip()
|
||||
if "Total Number of Cores" in line:
|
||||
info["GPU Cores"] = line.partition(":")[2].strip()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return info
|
||||
|
||||
def get_topology(self):
|
||||
return {}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if is_cuda_v2():
|
||||
env = GPUEnv()
|
||||
@@ -521,4 +597,6 @@ if __name__ == "__main__":
|
||||
env = NPUEnv()
|
||||
elif is_musa():
|
||||
env = MUSAEnv()
|
||||
elif is_mps():
|
||||
env = MPSEnv()
|
||||
env.check_env()
|
||||
|
||||
Reference in New Issue
Block a user