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 1178681f1533b..a409c49b5dc00 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