From b4b9813b5e2076d510ae518252f64064e6646a3e Mon Sep 17 00:00:00 2001 From: Zhang Jason Date: Wed, 6 Aug 2025 23:58:38 +0800 Subject: [PATCH] add the codes to check AMD Instinct GPU number (#22367) Signed-off-by: Zhang Jason --- .../disagg_prefill_lmcache_v1/disagg_example_nixl.sh | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh b/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh index 1178681f1533..a409c49b5dc0 100644 --- a/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh +++ b/examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh @@ -21,8 +21,14 @@ check_hf_token() { } check_num_gpus() { - # can you check if the number of GPUs are >=2 via nvidia-smi? - num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l) + # can you check if the number of GPUs are >=2 via nvidia-smi/rocm-smi? + which rocm-smi > /dev/null 2>&1 + if [ $? -ne 0 ]; then + num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l) + else + num_gpus=$(rocm-smi --showid | grep Instinct | wc -l) + fi + if [ "$num_gpus" -lt 2 ]; then echo "You need at least 2 GPUs to run disaggregated prefill." exit 1