From ab46acdbdadd1ea182af4bbc87882638e5dfdcf2 Mon Sep 17 00:00:00 2001
From: DontEatOreo <57304299+DontEatOreo@users.noreply.github.com>
Date: Mon, 8 Jul 2024 15:52:28 +0300
Subject: [PATCH] koboldcpp: fix CUDA
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Fix CUDA lib linkage
- Add new optional option `arches`

Refs: https://github.com/NixOS/nixpkgs/pull/314450#issuecomment-2170543384
Refs: https://github.com/NixOS/nixpkgs/pull/314450#issuecomment-2208635912
Refs: https://github.com/NixOS/nixpkgs/pull/325540#discussion_r1669900645

Co-authored-by: 百地 希留耶 <65301509+KiruyaMomochi@users.noreply.github.com>
Co-authored-by: mirrexagon <mirrexagon@mirrexagon.com>
Co-authored-by: SomeoneSerge <sergei.kozlukov@aalto.fi>
---
 pkgs/by-name/ko/koboldcpp/package.nix | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/pkgs/by-name/ko/koboldcpp/package.nix b/pkgs/by-name/ko/koboldcpp/package.nix
index 36ecdab4eac1..ba95550b08bf 100644
--- a/pkgs/by-name/ko/koboldcpp/package.nix
+++ b/pkgs/by-name/ko/koboldcpp/package.nix
@@ -6,6 +6,7 @@
   gitUpdater,
   python3Packages,
   tk,
+  addDriverRunpath,
 
   darwin,
 
@@ -28,6 +29,10 @@
 
   metalSupport ? stdenv.isDarwin && stdenv.isAarch64,
 
+  # You can find a full list here: https://arnon.dk/matching-sm-architectures-arch-and-gencode-for-various-nvidia-cards/
+  # For example if you're on an GTX 1080 that means you're using "Pascal" and you need to pass "sm_60"
+  arches ? cudaPackages.cudaFlags.arches or [ ],
+
   # You can find list of x86_64 options here: https://gcc.gnu.org/onlinedocs/gcc/x86-Options.html
   # For ARM here: https://gcc.gnu.org/onlinedocs/gcc/ARM-Options.html
   # If you set "march" to "native", specify "mtune" as well; otherwise, it will be set to "generic". (credit to: https://lemire.me/blog/2018/07/25/it-is-more-complicated-than-i-thought-mtune-march-in-gcc/)
@@ -39,6 +44,10 @@
 let
   makeBool = option: bool: (if bool then "${option}=1" else "");
 
+  makeWrapperArgs = lib.optionalString config.cudaSupport ''
+    --prefix LD_LIBRARY_PATH: "${lib.makeLibraryPath [ addDriverRunpath.driverLink ]}"
+  '';
+
   effectiveStdenv = if cublasSupport then cudaPackages.backendStdenv else stdenv;
 in
 effectiveStdenv.mkDerivation (finalAttrs: {
@@ -114,6 +123,7 @@ effectiveStdenv.mkDerivation (finalAttrs: {
     (makeBool "LLAMA_CLBLAST" clblastSupport)
     (makeBool "LLAMA_VULKAN" vulkanSupport)
     (makeBool "LLAMA_METAL" metalSupport)
+    (lib.optionalString cublasSupport "CUDA_DOCKER_ARCH=sm_${builtins.head arches}")
   ];
 
   installPhase = ''
@@ -136,7 +146,7 @@ effectiveStdenv.mkDerivation (finalAttrs: {
   postFixup = ''
     wrapPythonProgramsIn "$out/bin" "$pythonPath"
     makeWrapper "$out/bin/koboldcpp.unwrapped" "$out/bin/koboldcpp" \
-    --prefix PATH ${lib.makeBinPath [ tk ]}
+      --prefix PATH ${lib.makeBinPath [ tk ]} ${makeWrapperArgs}
   '';
 
   passthru.updateScript = gitUpdater { rev-prefix = "v"; };