fix(studio): opt-in source-build GPU smoke validation (#7322)

* fix(studio): opt-in source-build GPU smoke validation (#5854)

Gap 1 (empty CUDA arch -> CPU) already landed in #6481. Wire gap 2: after a
GPU source build, optionally run the same staged llama-server smoke test as
the prebuilt path, then CPU-fallback on failure. Gated by
UNSLOTH_LLAMA_STAGED_VALIDATION (default off) to avoid Blackwell JIT stalls.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(install): normalize staged validation env in setup.sh (#7322)

Strip and lowercase UNSLOTH_LLAMA_STAGED_VALIDATION before the shell
gate so values like True and surrounding whitespace match the Python
staged_validation_enabled() helper.

* Rebuild visual server after staged-validation CPU fallback (#5854)

Mirror the primary source-build path by best-effort building
llama-diffusion-gemma-visual-server after smoke-failure CPU fallback.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Souravrajvi0

.github/workflows/studio-backend-ci.yml

@@ -228,6 +228,7 @@ jobs:
               tests/sh/test_system_node_readonly.sh \
               tests/sh/test_nvcc_meets_llama_minimum.sh \
               tests/sh/test_resolve_cuda_archs.sh \
+              tests/sh/test_staged_validation_enabled.sh \
               tests/sh/test_tauri_install_exit_order.sh \
               tests/sh/test_torch_constraint.sh \
               tests/sh/test_torch_flavor.sh \

studio/install_llama_prebuilt.py

@@ -186,8 +186,24 @@ def env_int(
 # in validate_prebuilt_choice. Disabled for now: the llama-server GPU forward pass
 # JIT-compiles CUDA kernels on first load and stalls every install and update by
 # minutes on Blackwell (sm_100). The check and the source-build fallback it triggers
-# are kept intact -- set this to True to re-enable them.
+# are kept intact -- set this to True, or set UNSLOTH_LLAMA_STAGED_VALIDATION=1, to
+# re-enable them (#5854 gap 2).
 _RUN_STAGED_PREBUILT_VALIDATION = False
+
+
+def staged_validation_enabled() -> bool:
+    """True when the expensive llama-server GPU smoke test should run.
+
+    Default off (Blackwell CUDA JIT stalls installs). Opt in via the module
+    constant or ``UNSLOTH_LLAMA_STAGED_VALIDATION`` (1/true/yes/on). Used by both
+    the prebuilt path and setup.sh's source-build post-check (#5854).
+    """
+    if _RUN_STAGED_PREBUILT_VALIDATION:
+        return True
+    raw = os.environ.get("UNSLOTH_LLAMA_STAGED_VALIDATION", "").strip().lower()
+    return raw in ("1", "true", "yes", "on")
+
+
 INSTALL_LOCK_TIMEOUT_SECONDS = 300
 INSTALL_STAGING_ROOT_NAME = ".staging"
 GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"}
@@ -5868,9 +5884,10 @@ def validate_prebuilt_choice(
     # so they are always validated. For an approved bundle the sha256 manifest
     # already proves integrity, so its runtime smoke test -- a cold CUDA-JIT pass
     # costing minutes on Blackwell sm_100 -- is gated behind
-    # _RUN_STAGED_PREBUILT_VALIDATION, disabled for now. The check and the
-    # source-build fallback it triggers are kept intact; flip the flag to restore it.
-    if choice.expected_sha256 is None or _RUN_STAGED_PREBUILT_VALIDATION:
+    # staged_validation_enabled() (constant or UNSLOTH_LLAMA_STAGED_VALIDATION),
+    # disabled for now. The check and the source-build fallback it triggers are
+    # kept intact; flip the flag / env to restore it (#5854).
+    if choice.expected_sha256 is None or staged_validation_enabled():
         validate_quantize(
             quantize_path,
             probe_path,
@@ -5891,6 +5908,49 @@ def validate_prebuilt_choice(
     return server_path, quantize_path
 
 
+def validate_existing_install(
+    install_dir: Path,
+    *,
+    install_kind: str | None = None,
+    host: HostInfo | None = None,
+) -> None:
+    """Run the staged smoke test against an already-built llama.cpp tree (#5854).
+
+    Used by setup.sh after a GPU source build when ``UNSLOTH_LLAMA_STAGED_VALIDATION``
+    is set. Raises ``PrebuiltFallback`` on failure so the caller can retry CPU.
+    """
+    host = host or detect_host()
+    bin_dir = install_dir / "build" / "bin"
+    server_name = "llama-server.exe" if host.is_windows else "llama-server"
+    quantize_name = "llama-quantize.exe" if host.is_windows else "llama-quantize"
+    server_path = bin_dir / server_name
+    quantize_path = bin_dir / quantize_name
+    if not server_path.is_file():
+        raise PrebuiltFallback(f"llama-server not found at {server_path}")
+
+    with tempfile.TemporaryDirectory(prefix = "unsloth-llama-source-validate-") as tmp:
+        work_dir = Path(tmp)
+        probe_path = work_dir / "stories260K.gguf"
+        quantized_path = work_dir / "stories260K-q4.gguf"
+        download_validation_model(probe_path, validation_model_cache_path(install_dir))
+        if quantize_path.is_file():
+            validate_quantize(
+                quantize_path,
+                probe_path,
+                quantized_path,
+                install_dir,
+                host,
+            )
+        validate_server(
+            server_path,
+            probe_path,
+            host,
+            install_dir,
+            install_kind = install_kind,
+        )
+    log(f"staged source-build validation succeeded for {install_dir}")
+
+
 def validate_prebuilt_attempts(
     attempts: Iterable[AssetChoice],
     host: HostInfo,
@@ -6345,6 +6405,24 @@ def parse_args() -> argparse.Namespace:
             "fork). Use --output-format json."
         ),
     )
+    resolve_group.add_argument(
+        "--validate-install",
+        metavar = "DIR",
+        help = (
+            "Run the staged llama-server smoke test against an existing build "
+            "tree (setup.sh source-build post-check, #5854). Exit 2 on failure. "
+            "Normally gated by UNSLOTH_LLAMA_STAGED_VALIDATION; this flag always "
+            "runs the check."
+        ),
+    )
+    parser.add_argument(
+        "--install-kind",
+        default = None,
+        help = (
+            "Install kind for --validate-install GPU offload (e.g. linux-cuda, "
+            "linux-rocm, macos-arm64). When omitted, host detection decides."
+        ),
+    )
     parser.add_argument(
         "--output-format",
         choices = ("plain", "json"),
@@ -6381,6 +6459,17 @@ def emit_resolver_output(payload: dict[str, Any], *, output_format: str) -> None
 
 def main() -> int:
     args = parse_args()
+    if args.validate_install is not None:
+        try:
+            validate_existing_install(
+                Path(args.validate_install),
+                install_kind = args.install_kind,
+            )
+        except PrebuiltFallback as exc:
+            print(str(exc), file = sys.stderr)
+            raise SystemExit(EXIT_FALLBACK) from exc
+        return EXIT_SUCCESS
+
     if args.resolve_llama_tag is not None:
         resolved = resolve_requested_llama_tag(
             args.resolve_llama_tag,

studio/setup.sh

@@ -251,6 +251,38 @@ _resolve_cuda_archs() {
     printf '%s' "$_archs"
 }
 
+# Opt-in staged GPU smoke test after a source build (#5854 gap 2). Default off:
+# llama-server's first GPU forward pass JIT-compiles CUDA kernels and stalls
+# installs for minutes on Blackwell. Same env as install_llama_prebuilt.py.
+_staged_validation_enabled() {
+    local _raw="${UNSLOTH_LLAMA_STAGED_VALIDATION:-}"
+    # Match install_llama_prebuilt.py staged_validation_enabled(): strip + lowercase.
+    _raw="$(printf '%s' "$_raw" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' | tr '[:upper:]' '[:lower:]')"
+    case "$_raw" in
+        1|true|yes|on) return 0 ;;
+        *) return 1 ;;
+    esac
+}
+
+# Map the source-build GPU backend to install_llama_prebuilt --install-kind so
+# validate_server enables --n-gpu-layers for the right backends.
+_source_smoke_install_kind() {
+    if [ "${_TRY_METAL_CPU_FALLBACK:-false}" = true ]; then
+        printf '%s' "macos-arm64"
+        return 0
+    fi
+    case "${GPU_BACKEND:-}" in
+        cuda)
+            case "$(uname -m 2>/dev/null || true)" in
+                aarch64|arm64) printf '%s' "linux-arm64-cuda" ;;
+                *) printf '%s' "linux-cuda" ;;
+            esac
+            ;;
+        rocm) printf '%s' "linux-rocm" ;;
+        *) printf '%s' "" ;;
+    esac
+}
+
 # Run a GPU probe under a 10s timeout when `timeout` is available so a wedged
 # NVIDIA driver cannot hang setup; fall back to a bare call where it is not.
 _setup_run_smi() {
@@ -1900,6 +1932,37 @@ else
             run_quiet_no_exit "build diffusion visual server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-diffusion-gemma-visual-server -j"$NCPU" || true
         fi
 
+        # Opt-in post-build GPU smoke test (#5854 gap 2). Default off (Blackwell
+        # CUDA JIT stalls). On failure, reuse the CPU fallback path so the user
+        # still gets a working llama-server. Runs before the install swap.
+        if [ "$BUILD_OK" = true ] && _staged_validation_enabled; then
+            _FB_LABEL="$(_gpu_fallback_label)"
+            _SMOKE_KIND="$(_source_smoke_install_kind)"
+            if [ -n "$_FB_LABEL" ]; then
+                _SMOKE_CMD=(
+                    python "$SCRIPT_DIR/install_llama_prebuilt.py"
+                    --validate-install "$_BUILD_TMP"
+                )
+                [ -n "$_SMOKE_KIND" ] && _SMOKE_CMD+=(--install-kind "$_SMOKE_KIND")
+                if ! run_quiet_no_exit "validate source llama.cpp" "${_SMOKE_CMD[@]}"; then
+                    substep "$_FB_LABEL source build failed smoke test; retrying CPU build..." "$C_WARN"
+                    _TRY_METAL_CPU_FALLBACK=false
+                    rm -rf "$_BUILD_TMP/build"
+                    if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then
+                        _BUILD_DESC="building (CPU fallback after $_FB_LABEL smoke failed)"
+                        GPU_BACKEND=""
+                        run_quiet_no_exit "build llama-server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
+                        if [ "$BUILD_OK" = true ]; then
+                            run_quiet_no_exit "build llama-quantize (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-quantize -j"$NCPU" || true
+                            run_quiet_no_exit "build diffusion visual server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-diffusion-gemma-visual-server -j"$NCPU" || true
+                        fi
+                    else
+                        BUILD_OK=false
+                    fi
+                fi
+            fi
+        fi
+
         # Swap only after build succeeds -- preserves existing install on failure
         if [ "$BUILD_OK" = true ]; then
             _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install"

tests/run_all.sh

@@ -12,6 +12,7 @@ sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
 sh "$TESTS_DIR/sh/test_torch_constraint.sh"
 sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh"
 sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh"
+sh "$TESTS_DIR/sh/test_staged_validation_enabled.sh"
 sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh"
 sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh"
 sh "$TESTS_DIR/sh/test_torch_flavor.sh"

tests/sh/test_staged_validation_enabled.sh

@@ -0,0 +1,116 @@
+#!/bin/bash
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+# Unit tests for setup.sh staged-validation helpers (#5854 gap 2).
+# Opt-in GPU smoke after a source build; default off (Blackwell JIT stall).
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh"
+PASS=0
+FAIL=0
+
+_FUNC_FILE=$(mktemp)
+{
+    sed -n '/^_staged_validation_enabled()/,/^}/p' "$SETUP_SH"
+    sed -n '/^_source_smoke_install_kind()/,/^}/p' "$SETUP_SH"
+} > "$_FUNC_FILE"
+# shellcheck disable=SC1090
+. "$_FUNC_FILE"
+rm -f "$_FUNC_FILE"
+
+assert_eq() {
+    _label="$1"; _expected="$2"; _actual="$3"
+    if [ "$_actual" = "$_expected" ]; then
+        echo "  PASS: $_label"; PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1))
+    fi
+}
+
+assert_rc() {
+    _label="$1"; _expected="$2"
+    shift 2
+    set +e
+    "$@" >/dev/null 2>&1
+    _rc=$?
+    set -e
+    if [ "$_rc" -eq "$_expected" ]; then
+        echo "  PASS: $_label"; PASS=$((PASS + 1))
+    else
+        echo "  FAIL: $_label (expected rc $_expected, got $_rc)"; FAIL=$((FAIL + 1))
+    fi
+}
+
+echo "=== _staged_validation_enabled ==="
+unset UNSLOTH_LLAMA_STAGED_VALIDATION
+assert_rc "default off" 1 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION=0
+assert_rc "0 is off" 1 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION=1
+assert_rc "1 is on" 0 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION=true
+assert_rc "true is on" 0 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION=yes
+assert_rc "yes is on" 0 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION=on
+assert_rc "on is on" 0 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION=True
+assert_rc "True is on" 0 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION='  yes  '
+assert_rc "whitespace yes is on" 0 _staged_validation_enabled
+
+UNSLOTH_LLAMA_STAGED_VALIDATION=maybe
+assert_rc "maybe is off" 1 _staged_validation_enabled
+unset UNSLOTH_LLAMA_STAGED_VALIDATION
+
+echo "=== _source_smoke_install_kind ==="
+_TRY_METAL_CPU_FALLBACK=true
+GPU_BACKEND=""
+assert_eq "metal" "macos-arm64" "$(_source_smoke_install_kind)"
+
+_TRY_METAL_CPU_FALLBACK=false
+GPU_BACKEND=cuda
+_kind="$(_source_smoke_install_kind)"
+case "$(uname -m)" in
+    aarch64|arm64) assert_eq "cuda arm" "linux-arm64-cuda" "$_kind" ;;
+    *) assert_eq "cuda x86" "linux-cuda" "$_kind" ;;
+esac
+
+GPU_BACKEND=rocm
+assert_eq "rocm" "linux-rocm" "$(_source_smoke_install_kind)"
+
+GPU_BACKEND=""
+assert_eq "cpu empty" "" "$(_source_smoke_install_kind)"
+
+echo "=== setup.sh source smoke contract ==="
+assert_contains() {
+    _label="$1"; _hay="$2"; _needle="$3"
+    case "$_hay" in
+        *"$_needle"*) echo "  PASS: $_label"; PASS=$((PASS + 1)) ;;
+        *) echo "  FAIL: $_label (missing '$_needle')"; FAIL=$((FAIL + 1)) ;;
+    esac
+}
+_src=$(cat "$SETUP_SH")
+assert_contains "env gate present" "$_src" "UNSLOTH_LLAMA_STAGED_VALIDATION"
+assert_contains "calls validate-install" "$_src" "--validate-install"
+assert_contains "smoke fail retries CPU" "$_src" "source build failed smoke test; retrying CPU build"
+# Smoke must run before the install swap.
+_smoke_pos=$(printf '%s' "$_src" | awk '/validate source llama.cpp/{print NR; exit}')
+_swap_pos=$(printf '%s' "$_src" | awk '/mv "\$_BUILD_TMP" "\$LLAMA_CPP_DIR"/{print NR; exit}')
+if [ -n "$_smoke_pos" ] && [ -n "$_swap_pos" ] && [ "$_smoke_pos" -lt "$_swap_pos" ]; then
+    echo "  PASS: smoke before install swap"; PASS=$((PASS + 1))
+else
+    echo "  FAIL: smoke before install swap (smoke=$_smoke_pos swap=$_swap_pos)"; FAIL=$((FAIL + 1))
+fi
+
+echo ""
+echo "Results: $PASS passed, $FAIL failed"
+[ "$FAIL" -eq 0 ]