Fix token speed slow in machine has multi gpus (#1157)

* Update bat script windows choose GPU has highest ram to start nitro

* Update bash script for linux to choose gpu has highest vram

---------

Co-authored-by: Hien To <tominhhien97@gmail.com>

extensions/inference-nitro-extension/bin/linux-start.sh

@@ -1,5 +1,38 @@
 #!/bin/bash
 
-# Attempt to run the nitro_linux_amd64_cuda file and if it fails, run nitro_linux_amd64
+# Check if nvidia-smi exists and is executable
+if ! command -v nvidia-smi &> /dev/null; then
+    echo "nvidia-smi not found, proceeding with CPU version..."
+    cd linux-cpu
+    ./nitro "$@"
+    exit $?
+fi
+
+# Find the GPU with the highest VRAM
+readarray -t gpus < <(nvidia-smi --query-gpu=index,memory.total --format=csv,noheader,nounits)
+maxMemory=0
+selectedGpuId=0
+
+for gpu in "${gpus[@]}"; do
+    IFS=, read -ra gpuInfo <<< "$gpu"
+    gpuId=${gpuInfo[0]}
+    gpuMemory=${gpuInfo[1]}
+    if (( gpuMemory > maxMemory )); then
+        maxMemory=$gpuMemory
+        selectedGpuId=$gpuId
+    fi
+done
+
+echo "Selected GPU: $selectedGpuId"
+export CUDA_VISIBLE_DEVICES=$selectedGpuId
+
+# Attempt to run nitro_linux_amd64_cuda
 cd linux-cuda
-./nitro "$@" || (echo "nitro_linux_amd64_cuda encountered an error, attempting to run nitro_linux_amd64..." && cd ../linux-cpu && ./nitro "$@")
+if ./nitro "$@"; then
+    exit $?
+else
+    echo "nitro_linux_amd64_cuda encountered an error, attempting to run nitro_linux_amd64..."
+    cd ../linux-cpu
+    ./nitro "$@"
+    exit $?
+fi

extensions/inference-nitro-extension/bin/win-start.bat

@@ -1,12 +1,44 @@
 @echo off
+setlocal enabledelayedexpansion
+
+set "maxMemory=0"
+set "gpuId="
+
+rem check if nvidia-smi command exist or not
+where nvidia-smi >nul 2>&1
+if %errorlevel% neq 0 (
+    echo nvidia-smi not found, proceeding with CPU version...
+    cd win-cuda
+    goto RunCpuVersion
+)
+
+set "tempFile=%temp%\nvidia_smi_output.txt"
+nvidia-smi --query-gpu=index,memory.total --format=csv,noheader,nounits > "%tempFile%"
+
+for /f "usebackq tokens=1-2 delims=, " %%a in ("%tempFile%") do (
+    set /a memorySize=%%b
+    if !memorySize! gtr !maxMemory! (
+        set "maxMemory=!memorySize!"
+        set "gpuId=%%a"
+    )
+)
+
+rem Echo the selected GPU
+echo Selected GPU: !gpuId!
+
+rem Set the GPU with the highest VRAM as the visible CUDA device
+set CUDA_VISIBLE_DEVICES=!gpuId!
 
 rem Attempt to run nitro_windows_amd64_cuda.exe
 cd win-cuda
 nitro.exe %*
+if %errorlevel% neq 0 goto RunCpuVersion
+goto End
 
-rem Check the exit code of the previous command
-if %errorlevel% neq 0 (
-    echo nitro_windows_amd64_cuda.exe encountered an error, attempting to run nitro_windows_amd64.exe...
-    cd ..\win-cpu
-    nitro.exe %*
-)
+:RunCpuVersion
+rem Run nitro_windows_amd64.exe...
+cd ..\win-cpu
+nitro.exe %*
+
+:End
+endlocal