ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(studio): add pip-installed nvidia CUDA libs to LD_LIBRARY_PATH for llama-server (#4590)

Browse source 
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>
This commit is contained in:
Daniel Han 2026-03-25 06:24:40 -07:00 • committed by GitHub
commit ae2b1b97ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 40 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

40
studio/backend/core/inference/llama_cpp.py
View file

@ -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",