A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[ROCm] Tweak the benchmark script to run on ROCm (#14252)
Huy Do committed
Apr 3, 2025 at 17:12 UTC
092475f7386ea36c1aed17be65b872139931be98
1 file changed
+21
-6
.buildkite/nightly-benchmarks/scripts/run-performance-benchmarks.sh
+21
-6
@@ -10,15 +10,24 @@ set -x
10
set -o pipefail
11
12
check_gpus() {
13
- # check the number of GPUs and GPU type.
14
- declare -g gpu_count=$(nvidia-smi --list-gpus | wc -l)
13
+ if command -v nvidia-smi; then
14
+ # check the number of GPUs and GPU type.
15
+ declare -g gpu_count=$(nvidia-smi --list-gpus | wc -l)
16
+ elif command -v amd-smi; then
17
+ declare -g gpu_count=$(amd-smi list | grep 'GPU' | wc -l)
18
+ fi
19
+
20
if [[ $gpu_count -gt 0 ]]; then
21
echo "GPU found."
22
else
23
echo "Need at least 1 GPU to run benchmarking."
24
exit 1
25
fi
21
- declare -g gpu_type=$(nvidia-smi --query-gpu=name --format=csv,noheader | awk '{print $2}')
26
+ if command -v nvidia-smi; then
27
+ declare -g gpu_type=$(nvidia-smi --query-gpu=name --format=csv,noheader | awk '{print $2}')
28
+ elif command -v amd-smi; then
29
+ declare -g gpu_type=$(amd-smi static -g 0 -a | grep 'MARKET_NAME' | awk '{print $2}')
30
+ fi
31
echo "GPU type is $gpu_type"
32
}
33
@@ -90,9 +99,15 @@ kill_gpu_processes() {
99
100
101
# wait until GPU memory usage smaller than 1GB
93
- while [ "$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -n 1)" -ge 1000 ]; do
94
- sleep 1
95
- done
102
+ if command -v nvidia-smi; then
103
+ while [ "$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | head -n 1)" -ge 1000 ]; do
104
+ sleep 1
105
+ done
106
+ elif command -v amd-smi; then
107
+ while [ "$(amd-smi metric -g 0 | grep 'USED_VRAM' | awk '{print $2}')" -ge 1000 ]; do
108
+ sleep 1
109
+ done
110
+ fi
111
112
# remove vllm config file
113
rm -rf ~/.config/vllm