From a2c3f39e12ed7e10c2dca928630f7e9887dab2e4 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Tue, 17 Mar 2026 11:12:54 +0000
Subject: [PATCH] setup: support CUDA_ARCHS_EXTRA env for Docker builds

---
 studio/setup.sh | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/studio/setup.sh b/studio/setup.sh
index 94a593df03..ee95e03966 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -433,19 +433,34 @@ rm -rf "$LLAMA_CPP_DIR"
 
                 # Detect GPU compute capability and limit CUDA architectures
                 # Without this, cmake builds for ALL default archs (very slow)
-                CUDA_ARCHS="86"
+                # CUDA_ARCHS can be pre-set via env var (e.g. Docker ARG) to
+                # include targets not detectable at build time (no GPU access).
+                # We always include 86 (A10) and merge with detected GPUs.
+                _add_arch() {
+                    local _a="$1"
+                    case ";$CUDA_ARCHS;" in
+                        *";$_a;"*) ;;
+                        *) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_a" ;;
+                    esac
+                }
+                CUDA_ARCHS=""
+                # Merge any pre-set CUDA_ARCHS from environment
+                if [ -n "${CUDA_ARCHS_EXTRA:-}" ]; then
+                    IFS=';' read -ra _env_archs <<< "$CUDA_ARCHS_EXTRA"
+                    for _ea in "${_env_archs[@]}"; do
+                        _ea=$(echo "$_ea" | tr -d '[:space:]')
+                        [ -n "$_ea" ] && _add_arch "$_ea"
+                    done
+                fi
+                # Always include sm_86 (A10)
+                _add_arch "86"
                 if command -v nvidia-smi &>/dev/null; then
                     # Read all GPUs, deduplicate (handles mixed-GPU hosts)
                     _raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
                     while IFS= read -r _cap; do
                         _cap=$(echo "$_cap" | tr -d '[:space:]')
                         if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then
-                            _arch="${BASH_REMATCH[1]}${BASH_REMATCH[2]}"
-                            # Append if not already present
-                            case ";$CUDA_ARCHS;" in
-                                *";$_arch;"*) ;;
-                                *) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_arch" ;;
-                            esac
+                            _add_arch "${BASH_REMATCH[1]}${BASH_REMATCH[2]}"
                         fi
                     done <<< "$_raw_caps"
                 fi