docker: give llama.cpp its libcublas so GGUF stops running on the CPU

The portable llama.cpp bundle loads libggml-cuda.so with dlopen, links it
against libcublas, and does not ship libcublas. The CUDA runtime base
image only carries libcudart, and the only libcublas in the image is
torch's wheel copy under site-packages/nvidia/cublas/lib, which was not
on the loader path. So the CUDA backend failed to load, and llama.cpp
said nothing about it: `--list-devices` printed an empty list and every
GGUF request ran on the CPU.

Measured in the built image on a B200 with gemma-4-E2B-it UD-Q4_K_XL:
1.6 tok/s from llama-cli and 4.2 tok/s from llama-server. With the fix,
the same image and model report `CUDA0: NVIDIA B200` and run at 229 tok/s
and 193 tok/s. Studio's GGUF chat and the GGUF export path go through the
same bundle, so both were affected.

The venv loader config already existed for torchcodec, so cublas/lib
joins it there rather than on LD_LIBRARY_PATH: ld.so.conf.d is consulted
after DT_RUNPATH, which keeps llama.cpp resolving its own $ORIGIN libs
first. cu13/lib comes along for the arm64 bundle's layout.

A silent 140x slowdown deserves a build-time gate, so the layer after the
fetch runs ldd over libggml-cuda.so, installs the cublas major the bundle
actually asks for when it is missing, and fails the build on anything
still unresolved. The amd64 bundle wants libcublas.so.12 and torch
already provides it; the arm64 bundle is CUDA 13, and deriving the major
from ldd keeps that leg honest without hardcoding either. libcuda.so.1 is
exempt: nvidia-container-toolkit injects the driver stub at
`docker run --gpus`, so it is never resolvable at build time. ldd needs
no GPU, so the build stays host-independent.

tests/python/test_docker_llama_cuda_backend.py pins the loader entry, the
guard, the driver-stub exemption and the ordering.
This commit is contained in:
Daniel Han 2026-07-26 14:59:27 +00:00
commit 9e3f3671d0
2 changed files with 127 additions and 1 deletions

View file

@ -0,0 +1,89 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0
"""Regression guard for the llama.cpp CUDA backend inside the Docker image.
The portable llama.cpp bundle ships libggml-cuda.so and loads it with dlopen
(ggml_backend_dl), but the bundle does NOT carry the CUDA math libraries it
links against, and the CUDA runtime base image only carries libcudart. With no
libcublas on the loader path the backend fails to load SILENTLY: llama.cpp
prints nothing, `--list-devices` comes back empty and every GGUF request runs on
the CPU. Measured on a B200 with gemma-4-E2B UD-Q4_K_XL: 1.6 tok/s instead of
224 tok/s, a 140x regression that no functional test would have caught.
The Dockerfile therefore has to do two things, and these tests pin both:
* put torch's bundled libcublas on the loader path (ld.so.conf.d, not
LD_LIBRARY_PATH, so llama.cpp's own $ORIGIN libs keep winning);
* fail the build when any non-driver dependency of libggml-cuda.so is still
unresolved, so a CPU-only image can never be published again.
Static: parses the Dockerfile only. No docker, no GPU, no network.
"""
from __future__ import annotations
import re
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
DOCKERFILE = REPO_ROOT / "docker" / "Dockerfile"
@pytest.fixture(scope="module")
def dockerfile() -> str:
assert DOCKERFILE.is_file(), f"missing {DOCKERFILE}"
return DOCKERFILE.read_text()
def test_cublas_dir_is_registered_with_the_loader(dockerfile: str):
conf = re.search(
r"ld\.so\.conf\.d/zz-unsloth-venv\.conf", dockerfile,
)
assert conf, "the venv loader-config layer disappeared"
block = dockerfile[: conf.end()]
assert "$SP/nvidia/cublas/lib" in block, (
"libggml-cuda.so links against libcublas, which only exists in the venv's "
"wheel copy; without this entry the CUDA backend fails to dlopen and GGUF "
"silently runs on the CPU"
)
def test_loader_config_is_not_ld_library_path(dockerfile: str):
# LD_LIBRARY_PATH is consulted BEFORE DT_RUNPATH, so it would let the venv's
# copies shadow llama.cpp's own $ORIGIN libs. ld.so.conf.d is consulted after.
assert "ld.so.conf.d/zz-unsloth-venv.conf" in dockerfile
assert not re.search(
r"ENV\s+LD_LIBRARY_PATH=.*site-packages/nvidia", dockerfile,
), "the venv nvidia libs must not go on LD_LIBRARY_PATH"
def test_build_fails_on_an_unresolved_cuda_backend(dockerfile: str):
assert "libggml-cuda.so" in dockerfile, "the CUDA backend guard disappeared"
guard = dockerfile[dockerfile.index("CUDA_SO=") :]
assert "ldd" in guard, "the guard must inspect the backend's dependencies"
assert "not found" in guard
assert "exit 1" in guard, "an unresolved backend must fail the build"
# The driver stub is injected by nvidia-container-toolkit at `docker run
# --gpus`, so it is never resolvable inside the build and must be exempt.
assert re.search(r"grep -v .libcuda\\?\.so\\?\.1", guard), (
"libcuda.so.1 must be exempt from the guard or every build fails"
)
def test_guard_installs_the_matching_cublas_major(dockerfile: str):
# The amd64 bundle is CUDA 12 and torch already ships libcublas.so.12, but
# the arm64 bundle is CUDA 13. Deriving the major from ldd keeps the two
# legs correct without hardcoding either.
guard = dockerfile[dockerfile.index("CUDA_SO=") :]
assert "nvidia-cublas-cu${major}" in guard, (
"the guard must install the cublas major the bundle actually asks for"
)
assert "libcublas" in guard
def test_guard_runs_after_the_prebuilt_is_fetched(dockerfile: str):
fetch = dockerfile.index("fetch_llama_prebuilt.py")
guard = dockerfile.index("CUDA_SO=")
assert fetch < guard, "the guard can only inspect a bundle that already exists"