From ae2b1b97ba24b96b82ebab8a55facc5e27645ca2 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 25 Mar 2026 06:24:40 -0700
Subject: [PATCH] fix(studio): add pip-installed nvidia CUDA libs to
 LD_LIBRARY_PATH for llama-server (#4590)

The prebuilt llama.cpp binary (cuda13-newer) links against
libcudart.so.13 and libcublas.so.13. When torch is installed via pip,
these libraries live in the venv's site-packages under
nvidia/cu13/lib/, not in /usr/local/cuda/.

The existing LD_LIBRARY_PATH logic only searched /usr/local/cuda*
paths (which have CUDA 12.x), so the CUDA backend failed to load
silently and llama-server fell back to CPU -- even with -ngl -1.

This adds a glob scan of the venv's nvidia package directories
(cu*, cudnn, nvjitlink) to LD_LIBRARY_PATH before launching
llama-server, matching where pip puts the CUDA runtime.

Tested on Colab with RTX PRO 6000 Blackwell (CUDA 13.0, pip torch):
before -- 3 MiB GPU, 0% util, CPU inference
after  -- 13317 MiB GPU, 77% util, full GPU inference

Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
---
 studio/backend/core/inference/llama_cpp.py | 40 ++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 81a087341a..1d5643ac09 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -969,6 +969,46 @@ class LlamaCppBackend:
 
                 lib_dirs = [binary_dir]
                 _arch = platform.machine()  # x86_64, aarch64, etc.
+
+                # Pip-installed nvidia CUDA runtime libs (e.g. torch's
+                # bundled cuda-bindings).  The prebuilt llama.cpp binary
+                # links against libcudart.so.13 / libcublas.so.13 which
+                # live here, not in /usr/local/cuda.
+                import glob as _glob
+
+                for _nv_pattern in [
+                    os.path.join(
+                        sys.prefix,
+                        "lib",
+                        "python*",
+                        "site-packages",
+                        "nvidia",
+                        "cu*",
+                        "lib",
+                    ),
+                    os.path.join(
+                        sys.prefix,
+                        "lib",
+                        "python*",
+                        "site-packages",
+                        "nvidia",
+                        "cudnn",
+                        "lib",
+                    ),
+                    os.path.join(
+                        sys.prefix,
+                        "lib",
+                        "python*",
+                        "site-packages",
+                        "nvidia",
+                        "nvjitlink",
+                        "lib",
+                    ),
+                ]:
+                    for _nv_dir in _glob.glob(_nv_pattern):
+                        if os.path.isdir(_nv_dir):
+                            lib_dirs.append(_nv_dir)
+
                 for cuda_lib in [
                     "/usr/local/cuda/lib64",
                     f"/usr/local/cuda/targets/{_arch}-linux/lib",