Skip to content

Commit

Permalink
nix: add cuda to the flake
Browse files Browse the repository at this point in the history
  • Loading branch information
Green-Sky committed Sep 18, 2023
1 parent 7ddf185 commit 1b55cb9
Showing 1 changed file with 24 additions and 0 deletions.
24 changes: 24 additions & 0 deletions flake.nix
Original file line number Diff line number Diff line change
Expand Up @@ -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 ];
Expand Down

0 comments on commit 1b55cb9

Please sign in to comment.