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, 08:30 UTC
[AMD] [P/D] Compute num gpus for ROCm correctly in run_accuracy_test.sh (#18568)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
May 22, 2025 at 20:45 UTC
46791e1b4bad8698043eb632bb856306921f2334
1 file changed
+11
-2
tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh
+11
-2
@@ -13,6 +13,8 @@ NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-2} # Default to 2
13
# Find the git repository root directory
14
GIT_ROOT=$(git rev-parse --show-toplevel)
15
16
+SMI_BIN=$(which nvidia-smi || which rocm-smi)
17
+
18
# Trap the SIGINT signal (triggered by Ctrl+C)
19
trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
20
@@ -44,6 +46,13 @@ get_model_args() {
46
echo "$extra_args"
47
}
48
49
+get_num_gpus() {
50
+ if [[ "$SMI_BIN" == *"nvidia"* ]]; then
51
+ echo "$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)"
52
+ else
53
+ echo "$($SMI_BIN -l | grep GPU | wc -l)"
54
+ fi
55
+}
56
57
# Function to run tests for a specific model
58
run_tests_for_model() {
@@ -64,7 +73,7 @@ run_tests_for_model() {
73
# Start prefill instances
74
for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do
75
# Calculate GPU ID - we'll distribute across available GPUs
67
- GPU_ID=$((i % $(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)))
76
+ GPU_ID=$((i % $(get_num_gpus)))
77
# Calculate port number (base port + instance number)
78
PORT=$((8100 + i))
79
# Calculate side channel port
@@ -96,7 +105,7 @@ run_tests_for_model() {
105
# Start decode instances
106
for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do
107
# Calculate GPU ID - we'll distribute across available GPUs, starting from after prefill GPUs
99
- GPU_ID=$(((i + NUM_PREFILL_INSTANCES) % $(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)))
108
+ GPU_ID=$(((i + NUM_PREFILL_INSTANCES) % $(get_num_gpus)))
109
# Calculate port number (base port + instance number)
110
PORT=$((8200 + i))
111
# Calculate side channel port