Repository: ollama/ollama
Pull request: Detect very old CUDA GPUs and fall back to CPU

If we try to load the CUDA library on an old GPU, it panics and crashes the server.  This checks the compute capability before we load the library so we can gracefully fall back to CPU mode.

Prior to version 0.1.18, the fallback behavior resulted from the subprocess runner crashing.  Example from an old GPU:
```
ggml_init_cublas: found 1 CUDA devices:
  Device 0: NVIDIA GeForce GTX 765M, compute capability 3.0

cuBLAS error 3 at /go/src/github.com/jmorganca/ollama/llm/llama.cpp/gguf/ggml-cuda.cu:5552
current device: 0
2024/01/06 21:48:54 llama.go:320: llama runner exited with error: exit status 1
```
In 0.1.18 without this fix, the server crashes with a panic due to the assert in llama.cpp.

With this fix on the same system we now detect and fallback to CPU mode:
```
2024/01/06 21:52:17 shim_ext_server.go:142: Dynamic LLM variants [cuda rocm]
2024/01/06 21:52:17 gpu.go:37: Detecting GPU type
2024/01/06 21:52:17 gpu.go:56: Nvidia GPU detected
2024/01/06 21:52:17 gpu.go:89: CUDA GPU is too old. Falling back to CPU mode. Compute Capability detected: 3.0
2024/01/06 21:52:17 routes.go:953: no GPU detected
...
```

Example output on a newer supported GPU which correctly runs on the GPU:
```
2024/01/06 21:55:11 shim_ext_server.go:142: Dynamic LLM variants [cuda rocm]
2024/01/06 21:55:11 gpu.go:37: Detecting GPU type
2024/01/06 21:55:11 gpu.go:56: Nvidia GPU detected
2024/01/06 21:55:11 gpu.go:86: CUDA Compute Capability detected: 7.5
```

Changed files:
filename: gpu/gpu.go
patch: @@ -28,6 +28,9 @@ type handles struct {
 var gpuMutex sync.Mutex
 var gpuHandles *handles = nil
 
+// TODO verify this is the correct min version
+const CudaComputeMajorMin = 5
+
 // Note: gpuMutex must already be held
 func initGPUHandles() {
 	// TODO - if the ollama build is CPU only, don't do these checks as they're irrelevant and confusing
@@ -73,7 +76,18 @@ func GetGPUInfo() GpuInfo {
 			log.Printf("error looking up CUDA GPU memory: %s", C.GoString(memInfo.err))
 			C.free(unsafe.Pointer(memInfo.err))
 		} else {
-			resp.Library = "cuda"
+			// Verify minimum compute capability
+			var cc C.cuda_compute_capability_t
+			C.cuda_compute_capability(*gpuHandles.cuda, &cc)
+			if cc.err != nil {
+				log.Printf("error looking up CUDA GPU compute capability: %s", C.GoString(cc.err))
+				C.free(unsafe.Pointer(cc.err))
+			} else if cc.major >= CudaComputeMajorMin {
+				log.Printf("CUDA Compute Capability detected: %d.%d", cc.major, cc.minor)
+				resp.Library = "cuda"
+			} else {
+				log.Printf("CUDA GPU is too old. Falling back to CPU mode. Compute Capability detected: %d.%d", cc.major, cc.minor)
+			}
 		}
 	} else if gpuHandles.rocm != nil {
 		C.rocm_check_vram(*gpuHandles.rocm, &memInfo)

filename: gpu/gpu_info_cuda.c
patch: @@ -21,7 +21,7 @@ const char *cuda_lib_paths[] = {
 };
 #endif
 
-#define CUDA_LOOKUP_SIZE 5
+#define CUDA_LOOKUP_SIZE 6
 
 void cuda_init(cuda_init_resp_t *resp) {
   nvmlReturn_t ret;
@@ -39,6 +39,7 @@ void cuda_init(cuda_init_resp_t *resp) {
       {"nvmlDeviceGetHandleByIndex", (void *)&resp->ch.getHandle},
       {"nvmlDeviceGetMemoryInfo", (void *)&resp->ch.getMemInfo},
       {"nvmlDeviceGetCount_v2", (void *)&resp->ch.getCount},
+      {"nvmlDeviceGetCudaComputeCapability", (void *)&resp->ch.getComputeCapability},
   };
 
   for (i = 0; cuda_lib_paths[i] != NULL && resp->ch.handle == NULL; i++) {
@@ -123,4 +124,53 @@ void cuda_check_vram(cuda_handle_t h, mem_info_t *resp) {
     resp->free += memInfo.free;
   }
 }
+
+void cuda_compute_capability(cuda_handle_t h, cuda_compute_capability_t *resp) {
+  resp->err = NULL;
+  resp->major = 0;
+  resp->minor = 0;
+  nvmlDevice_t device;
+  int major = 0;
+  int minor = 0;
+  nvmlReturn_t ret;
+  const int buflen = 256;
+  char buf[buflen + 1];
+  int i;
+
+  if (h.handle == NULL) {
+    resp->err = strdup("nvml handle not initialized");
+    return;
+  }
+
+  unsigned int devices;
+  ret = (*h.getCount)(&devices);
+  if (ret != NVML_SUCCESS) {
+    snprintf(buf, buflen, "unable to get device count: %d", ret);
+    resp->err = strdup(buf);
+    return;
+  }
+
+  for (i = 0; i < devices; i++) {
+    ret = (*h.getHandle)(i, &device);
+    if (ret != NVML_SUCCESS) {
+      snprintf(buf, buflen, "unable to get device handle %d: %d", i, ret);
+      resp->err = strdup(buf);
+      return;
+    }
+
+    ret = (*h.getComputeCapability)(device, &major, &minor);
+    if (ret != NVML_SUCCESS) {
+      snprintf(buf, buflen, "device compute capability lookup failure %d: %d", i, ret);
+      resp->err = strdup(buf);
+      return;
+    }
+    // Report the lowest major.minor we detect as that limits our compatibility
+    if (resp->major == 0 || resp->major > major ) {
+      resp->major = major;
+      resp->minor = minor;
+    } else if ( resp->major == major && resp->minor > minor ) {
+      resp->minor = minor;
+    }
+  }
+}
 #endif  // __APPLE__
\ No newline at end of file

filename: gpu/gpu_info_cuda.h
patch: @@ -22,15 +22,23 @@ typedef struct cuda_handle {
   nvmlReturn_t (*getHandle)(unsigned int, nvmlDevice_t *);
   nvmlReturn_t (*getMemInfo)(nvmlDevice_t, nvmlMemory_t *);
   nvmlReturn_t (*getCount)(unsigned int *);
+  nvmlReturn_t (*getComputeCapability)(nvmlDevice_t, int* major, int* minor);
 } cuda_handle_t;
 
 typedef struct cuda_init_resp {
   char *err;  // If err is non-null handle is invalid
   cuda_handle_t ch;
 } cuda_init_resp_t;
 
+typedef struct cuda_compute_capability {
+  char *err;
+  int major;
+  int minor;
+} cuda_compute_capability_t;
+
 void cuda_init(cuda_init_resp_t *resp);
 void cuda_check_vram(cuda_handle_t ch, mem_info_t *resp);
+void cuda_compute_capability(cuda_handle_t ch, cuda_compute_capability_t *cc);
 
 #endif  // __GPU_INFO_CUDA_H__
 #endif  // __APPLE__
\ No newline at end of file