nix: respect config.cudaCapabilities

ggerganov · Dec 26, 2023 · ae6bebc · ae6bebc
1 parent 4522c47
commit ae6bebc
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/.devops/nix/package.nix b/.devops/nix/package.nix
@@ -150,6 +150,14 @@ effectiveStdenv.mkDerivation (
         (cmakeBool "LLAMA_HIPBLAS" useRocm)
         (cmakeBool "LLAMA_METAL" useMetalKit)
       ]
+      ++ optionals useCuda [
+        (
+          with cudaPackages.flags;
+          cmakeFeature "CMAKE_CUDA_ARCHITECTURES" (
+            builtins.concatStringsSep ";" (map dropDot cudaCapabilities)
+          )
+        )
+      ]
       ++ optionals useRocm [
         (cmakeFeature "CMAKE_C_COMPILER" "hipcc")
         (cmakeFeature "CMAKE_CXX_COMPILER" "hipcc")