scripts/test_gpu.py

"""Verify llama-cpp-python is installed with GPU (CUDA) support."""

import sys


def check_gpu_support():
    try:
        import llama_cpp
    except ImportError:
        print("FAIL: llama-cpp-python is not installed")
        sys.exit(1)

    print(f"llama-cpp-python version: {llama_cpp.__version__}")

    has_gpu = llama_cpp.llama_supports_gpu_offload()
    print(f"GPU offload supported: {has_gpu}")

    if not has_gpu:
        print("FAIL: llama-cpp-python was built WITHOUT GPU support")
        print("Reinstall with: CMAKE_ARGS=\"-DGGML_CUDA=on\" pip install llama-cpp-python --force-reinstall --no-cache-dir --no-binary llama-cpp-python")
        sys.exit(1)

    print("OK: GPU support confirmed")


if __name__ == "__main__":
    check_gpu_support()