From 681a91499010be819dd45a1390e668b0817e7338 Mon Sep 17 00:00:00 2001
From: Daniel Hiltgen <daniel@ollama.com>
Date: Sat, 20 Jan 2024 10:48:43 -0800
Subject: [PATCH] Add support for CUDA 5.2 cards

---
 gpu/gpu.go                | 7 ++++---
 llm/generate/gen_linux.sh | 2 +-
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/gpu/gpu.go b/gpu/gpu.go
index 67bd8352..06d022fe 100644
--- a/gpu/gpu.go
+++ b/gpu/gpu.go
@@ -29,8 +29,9 @@ type handles struct {
 var gpuMutex sync.Mutex
 var gpuHandles *handles = nil
 
-// With our current CUDA compile flags, 5.2 and older will not work properly
-const CudaComputeMajorMin = 6
+// With our current CUDA compile flags, older than 5.2 will not work properly
+// 5.0: CUDA error: no kernel image is available for execution on the device
+var CudaComputeMin = [2]C.int{5, 2}
 
 // Possible locations for the nvidia-ml library
 var CudaLinuxGlobs = []string{
@@ -133,7 +134,7 @@ func GetGPUInfo() GpuInfo {
 			if cc.err != nil {
 				slog.Info(fmt.Sprintf("error looking up CUDA GPU compute capability: %s", C.GoString(cc.err)))
 				C.free(unsafe.Pointer(cc.err))
-			} else if cc.major >= CudaComputeMajorMin {
+			} else if cc.major > CudaComputeMin[0] || (cc.major == CudaComputeMin[0] && cc.minor >= CudaComputeMin[1]) {
 				slog.Info(fmt.Sprintf("CUDA Compute Capability detected: %d.%d", cc.major, cc.minor))
 				resp.Library = "cuda"
 			} else {
diff --git a/llm/generate/gen_linux.sh b/llm/generate/gen_linux.sh
index 0fcf1356..f00d717c 100755
--- a/llm/generate/gen_linux.sh
+++ b/llm/generate/gen_linux.sh
@@ -125,7 +125,7 @@ if [ -d "${CUDA_LIB_DIR}" ]; then
     if [ -n "${CUDA_MAJOR}" ]; then
         CUDA_VARIANT=_v${CUDA_MAJOR}
     fi
-    CMAKE_DEFS="-DLLAMA_CUBLAS=on ${COMMON_CMAKE_DEFS} ${CMAKE_DEFS}"
+    CMAKE_DEFS="-DLLAMA_CUBLAS=on -DLLAMA_CUDA_FORCE_MMQ=on ${COMMON_CMAKE_DEFS} ${CMAKE_DEFS}"
     BUILD_DIR="${LLAMACPP_DIR}/build/linux/${ARCH}/cuda${CUDA_VARIANT}"
     EXTRA_LIBS="-L${CUDA_LIB_DIR} -lcudart -lcublas -lcublasLt -lcuda"
     build