EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Use rocminfo instead of rocm-smi for more OS/WSL support (#2310)

Browse Source
This commit is contained in:
HAI
2024-12-02 00:15:45 -08:00
committed by GitHub
parent 3c79ad35ca
commit c54bda300a
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/utils.py
View File

@@ -903,7 +903,9 @@ def get_amdgpu_memory_capacity():
try: try:
# Run rocm-smi and capture the output # Run rocm-smi and capture the output
result = subprocess.run( result = subprocess.run(
["rocm-smi --showmeminfo vram | grep 'Total Memory' | awk '{print $NF}'"], [
"rocminfo | grep 'gfx94' -A 100 | grep 'Pool 1' -A 5 | grep 'Size:' | awk '{print $2}'"
],
stdout=subprocess.PIPE, stdout=subprocess.PIPE,
stderr=subprocess.PIPE, stderr=subprocess.PIPE,
shell=True, shell=True,
@@ -914,9 +916,8 @@ def get_amdgpu_memory_capacity():
# Parse the output to extract memory values in MiB # Parse the output to extract memory values in MiB
memory_values = [ memory_values = [
float(mem) / 1024 / 1024 float(mem.split("(")[0].strip()) / 1024
for mem in result.stdout.strip().split("\n") for mem in result.stdout.strip().split("\n")
if re.match(r"^\d+(\.\d+)?$", mem.strip())
] ]
if not memory_values: if not memory_values: