[MPS] Support sglang.check_env (#20753)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
This commit is contained in:
R0CKSTAR
2026-03-25 20:59:25 -07:00
committed by GitHub
parent d9e96153de
commit 02521420b3
+79 -1
View File
@@ -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()