From 1b55cb99742965669991761f828a1227a4c5bd15 Mon Sep 17 00:00:00 2001
From: Green Sky <green@g-s.xyz>
Date: Fri, 15 Sep 2023 22:42:00 +0200
Subject: [PATCH] nix: add cuda to the flake

---
 flake.nix | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/flake.nix b/flake.nix
index b0fb8642cb6188..673f7021771c05 100644
--- a/flake.nix
+++ b/flake.nix
@@ -69,6 +69,30 @@
             "-DLLAMA_CLBLAST=ON"
           ];
         };
+        packages.cuda = pkgs.stdenv.mkDerivation {
+          inherit name src meta postPatch nativeBuildInputs postInstall;
+          buildInputs = with pkgs; buildInputs ++ [ cudaPackages.cudatoolkit ];
+
+          # HACK(Green-Sky): nix currently has issues with cmake findcudatoolkit, so we force gnu make
+          # see https://github.com/NixOS/nixpkgs/issues/224291
+          dontUseCmakeConfigure = true;
+          dontUseNinjaBuild = true;
+
+          buildFlags = [ "LLAMA_CUBLAS=1" ];
+          installPhase = ''
+            mkdir -p $out/bin
+
+            # TODO(Green-Sky): add install target to Makefile, or wait for cmake support
+            mv main $out/bin/
+            mv server $out/bin/
+            mv speculative $out/bin/
+            mv perplexity $out/bin/
+            mv embedding $out/bin/
+            mv quantize $out/bin/
+            mv llama-bench $out/bin/
+            mv train-text-from-scratch $out/bin/
+          '';
+        };
         packages.rocm = pkgs.stdenv.mkDerivation {
           inherit name src meta postPatch nativeBuildInputs postInstall;
           buildInputs = with pkgs; buildInputs ++ [ hip hipblas rocblas ];