Skip to content

Commit 7c93fee

Browse files
zhangnjunoamgat
authored andcommitted
add the codes to check AMD Instinct GPU number (vllm-project#22367)
Signed-off-by: Zhang Jason <[email protected]> Signed-off-by: Noam Gat <[email protected]>
1 parent 48f0007 commit 7c93fee

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,14 @@ check_hf_token() {
2121
}
2222

2323
check_num_gpus() {
24-
# can you check if the number of GPUs are >=2 via nvidia-smi?
25-
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
24+
# can you check if the number of GPUs are >=2 via nvidia-smi/rocm-smi?
25+
which rocm-smi > /dev/null 2>&1
26+
if [ $? -ne 0 ]; then
27+
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
28+
else
29+
num_gpus=$(rocm-smi --showid | grep Instinct | wc -l)
30+
fi
31+
2632
if [ "$num_gpus" -lt 2 ]; then
2733
echo "You need at least 2 GPUs to run disaggregated prefill."
2834
exit 1

0 commit comments

Comments
 (0)