Skip to content

Add ROCm (AMD GPU) support to studio setup #4585

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
danielhanchen merged 3 commits into from
Mar 25, 2026
Merged

Add ROCm (AMD GPU) support to studio setup #4585

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6fd4961
Add support for ROCm in studio setup
edamamez Mar 18, 2026
24e22ca
Fix ROCm detection bugs: ROCM_PATH resolution, CUDA guard, compiler s…
danielhanchen Mar 25, 2026
5aa79ef
Fix gfx regex to match gfx90a (MI210/MI250/MI250X)
danielhanchen Mar 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 68 additions & 1 deletion studio/setup.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -449,17 +449,40 @@ rm -rf "$LLAMA_CPP_DIR"
echo " Using ccache for faster compilation"
fi

# Detect CUDA: check nvcc on PATH, then common install locations
# Detect GPU backend: CUDA (NVIDIA) or ROCm (AMD)
GPU_BACKEND=""

# Check for CUDA: check nvcc on PATH, then common install locations
NVCC_PATH=""
if command -v nvcc &>/dev/null; then
NVCC_PATH="$(command -v nvcc)"
GPU_BACKEND="cuda"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
GPU_BACKEND="cuda"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
GPU_BACKEND="cuda"
Comment on lines +459 to +468
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The GPU_BACKEND variable is assigned redundantly within each if/elif branch of the CUDA detection logic. It's more maintainable to set GPU_BACKEND once after NVCC_PATH has been successfully determined.

Suggested change
GPU_BACKEND="cuda"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
GPU_BACKEND="cuda"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
GPU_BACKEND="cuda"
NVCC_PATH="$(command -v nvcc)"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
fi
if [ -n "$NVCC_PATH" ]; then
GPU_BACKEND="cuda"
fi

fi
Comment on lines 457 to +469
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To improve maintainability and reduce code duplication, you can set GPU_BACKEND once after NVCC_PATH is found, instead of setting it in each branch of the conditional.

Suggested change
if command -v nvcc &>/dev/null; then
NVCC_PATH="$(command -v nvcc)"
GPU_BACKEND="cuda"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
GPU_BACKEND="cuda"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
GPU_BACKEND="cuda"
fi
if command -v nvcc &>/dev/null; then
NVCC_PATH="$(command -v nvcc)"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
fi
if [ -n "$NVCC_PATH" ]; then
GPU_BACKEND="cuda"
fi


# Check for ROCm (AMD) only if CUDA was not already selected
ROCM_HIPCC=""
if [ -z "$GPU_BACKEND" ]; then
if command -v hipcc &>/dev/null; then
ROCM_HIPCC="$(command -v hipcc)"
GPU_BACKEND="rocm"
elif [ -x /opt/rocm/bin/hipcc ]; then
ROCM_HIPCC="/opt/rocm/bin/hipcc"
export PATH="/opt/rocm/bin:$PATH"
GPU_BACKEND="rocm"
elif ls /opt/rocm-*/bin/hipcc &>/dev/null 2>&1; then
ROCM_HIPCC="$(ls -d /opt/rocm-*/bin/hipcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$ROCM_HIPCC"):$PATH"
GPU_BACKEND="rocm"
Comment on lines +475 to +484
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to the CUDA detection, the GPU_BACKEND variable is assigned redundantly within each if/elif branch of the ROCm detection logic. It's more maintainable to set GPU_BACKEND once after ROCM_HIPCC has been successfully determined.

Suggested change
ROCM_HIPCC="$(command -v hipcc)"
GPU_BACKEND="rocm"
elif [ -x /opt/rocm/bin/hipcc ]; then
ROCM_HIPCC="/opt/rocm/bin/hipcc"
export PATH="/opt/rocm/bin:$PATH"
GPU_BACKEND="rocm"
elif ls /opt/rocm-*/bin/hipcc &>/dev/null 2>&1; then
ROCM_HIPCC="$(ls -d /opt/rocm-*/bin/hipcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$ROCM_HIPCC"):$PATH"
GPU_BACKEND="rocm"
ROCM_HIPCC="$(command -v hipcc)"
elif [ -x /opt/rocm/bin/hipcc ]; then
ROCM_HIPCC="/opt/rocm/bin/hipcc"
export PATH="/opt/rocm/bin:$PATH"
elif ls /opt/rocm-*/bin/hipcc &>/dev/null 2>&1; then
ROCM_HIPCC="$(ls -d /opt/rocm-*/bin/hipcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$ROCM_HIPCC"):$PATH"
fi
if [ -n "$ROCM_HIPCC" ]; then
GPU_BACKEND="rocm"
fi

fi
fi
Comment on lines +473 to 486
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to the CUDA detection block, you can refactor this to set GPU_BACKEND only once after ROCM_HIPCC is found. This avoids code duplication and improves readability.

Suggested change
if [ -z "$GPU_BACKEND" ]; then
if command -v hipcc &>/dev/null; then
ROCM_HIPCC="$(command -v hipcc)"
GPU_BACKEND="rocm"
elif [ -x /opt/rocm/bin/hipcc ]; then
ROCM_HIPCC="/opt/rocm/bin/hipcc"
export PATH="/opt/rocm/bin:$PATH"
GPU_BACKEND="rocm"
elif ls /opt/rocm-*/bin/hipcc &>/dev/null 2>&1; then
ROCM_HIPCC="$(ls -d /opt/rocm-*/bin/hipcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$ROCM_HIPCC"):$PATH"
GPU_BACKEND="rocm"
fi
fi
if [ -z "$GPU_BACKEND" ]; then
if command -v hipcc &>/dev/null; then
ROCM_HIPCC="$(command -v hipcc)"
elif [ -x /opt/rocm/bin/hipcc ]; then
ROCM_HIPCC="/opt/rocm/bin/hipcc"
export PATH="/opt/rocm/bin:$PATH"
elif ls /opt/rocm-*/bin/hipcc &>/dev/null 2>&1; then
ROCM_HIPCC="$(ls -d /opt/rocm-*/bin/hipcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$ROCM_HIPCC"):$PATH"
fi
if [ -n "$ROCM_HIPCC" ]; then
GPU_BACKEND="rocm"
fi
fi


if [ -n "$NVCC_PATH" ]; then
Expand Down Expand Up @@ -494,9 +517,53 @@ rm -rf "$LLAMA_CPP_DIR"

# Multi-threaded nvcc compilation (uses all CPU cores per .cu file)
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_FLAGS=--threads=0"
elif [ "$GPU_BACKEND" = "rocm" ]; then
# Resolve hipcc symlinks to find the real ROCm root
_HIPCC_REAL="$(readlink -f "$ROCM_HIPCC" 2>/dev/null || printf '%s' "$ROCM_HIPCC")"
ROCM_ROOT=""
if command -v hipconfig &>/dev/null; then
ROCM_ROOT="$(hipconfig -R 2>/dev/null || true)"
fi
if [ -z "$ROCM_ROOT" ]; then
ROCM_ROOT="$(cd "$(dirname "$_HIPCC_REAL")/.." 2>/dev/null && pwd)"
fi

echo " Building with ROCm support (AMD GPU, hipcc: $_HIPCC_REAL)..."
CMAKE_ARGS="$CMAKE_ARGS -DGGML_HIP=ON"
export ROCM_PATH="$ROCM_ROOT"
export HIP_PATH="$ROCM_ROOT"

# Use upstream-recommended HIP compiler (not legacy hipcc-as-CXX)
if command -v hipconfig &>/dev/null; then
_HIP_CLANG_DIR="$(hipconfig -l 2>/dev/null || true)"
[ -n "$_HIP_CLANG_DIR" ] && export HIPCXX="$_HIP_CLANG_DIR/clang"
fi

# Detect AMD GPU architecture (gfx target)
GPU_TARGETS=""
if command -v rocminfo &>/dev/null; then
_gfx_list=$(rocminfo 2>/dev/null | grep -oE 'gfx[0-9]{2,4}[a-z]?' | sort -u || true)
_valid_gfx=""
for _gfx in $_gfx_list; do
if [[ "$_gfx" =~ ^gfx[0-9]{2,4}[a-z]?$ ]]; then
_valid_gfx="${_valid_gfx}${_valid_gfx:+;}$_gfx"
fi
done
Comment on lines +547 to +551
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The grep -oE 'gfx[0-9]{3,4}[a-z]?' command already filters for valid gfx patterns. The subsequent if [[ "$_gfx" =~ ^gfx[0-9]{3,4}[a-z]?$ ]] check inside the loop is redundant. You can directly append _gfx to _valid_gfx since it's already validated by grep.

Suggested change
for _gfx in $_gfx_list; do
if [[ "$_gfx" =~ ^gfx[0-9]{3,4}[a-z]?$ ]]; then
_valid_gfx="${_valid_gfx}${_valid_gfx:+;}$_gfx"
fi
done
_gfx_list=$(rocminfo 2>/dev/null | grep -oE 'gfx[0-9]{3,4}[a-z]?' | sort -u || true)
_valid_gfx=""
for _gfx in $_gfx_list; do
_valid_gfx="${_valid_gfx}${_valid_gfx:+;}$_gfx"
done
[ -n "$_valid_gfx" ] && GPU_TARGETS="$_valid_gfx"

[ -n "$_valid_gfx" ] && GPU_TARGETS="$_valid_gfx"
Comment on lines +546 to +552
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The for loop and if condition here are redundant because the grep command on the preceding line already filters for the correct gfx pattern. You can simplify this block to directly convert the newline-separated list from grep into a semicolon-separated string for GPU_TARGETS.

Suggested change
_valid_gfx=""
for _gfx in $_gfx_list; do
if [[ "$_gfx" =~ ^gfx[0-9]{3,4}[a-z]?$ ]]; then
_valid_gfx="${_valid_gfx}${_valid_gfx:+;}$_gfx"
fi
done
[ -n "$_valid_gfx" ] && GPU_TARGETS="$_valid_gfx"
if [ -n "$_gfx_list" ]; then
GPU_TARGETS=$(echo "$_gfx_list" | paste -sd ';')
fi

Comment on lines +546 to +552
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This block for building the GPU_TARGETS string can be simplified. The if condition inside the loop is redundant because grep -oE already ensures that all items in _gfx_list match the required pattern. Additionally, you can build the GPU_TARGETS string directly instead of using an intermediate _valid_gfx variable, since GPU_TARGETS is already initialized to an empty string on line 543.

Suggested change
_valid_gfx=""
for _gfx in $_gfx_list; do
if [[ "$_gfx" =~ ^gfx[0-9]{2,4}[a-z]?$ ]]; then
_valid_gfx="${_valid_gfx}${_valid_gfx:+;}$_gfx"
fi
done
[ -n "$_valid_gfx" ] && GPU_TARGETS="$_valid_gfx"
for _gfx in $_gfx_list; do
GPU_TARGETS="${GPU_TARGETS}${GPU_TARGETS:+;}$_gfx"
done

fi

if [ -n "$GPU_TARGETS" ]; then
echo " AMD GPU architectures: ${GPU_TARGETS//;/, } -- limiting build to detected targets"
CMAKE_ARGS="$CMAKE_ARGS -DGPU_TARGETS=${GPU_TARGETS}"
else
echo " Could not detect AMD GPU arch -- building for default targets (cmake will auto-detect)"
fi
elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then
echo " CUDA driver detected but nvcc not found — building CPU-only"
echo " To enable GPU: install cuda-toolkit or add nvcc to PATH"
elif [ -d /opt/rocm ] || command -v rocm-smi &>/dev/null; then
echo " ROCm driver detected but hipcc not found — building CPU-only"
echo " To enable GPU: install rocm-dev or add hipcc to PATH"
else
echo " Building CPU-only (no CUDA detected)..."
fi
Expand Down
Loading