python3Packages.xformers: 0.03 -> 0.0.23.post1, add cuda support

pkgs/development/python-modules/xformers/0001-fix-allow-building-without-git.patch

@@ -0,0 +1,25 @@
+From 4c2b65c47d328c2f20cc74adcec2286fee6cb5de Mon Sep 17 00:00:00 2001
+From: Yaroslav Bolyukin <iam@lach.pw>
+Date: Tue, 30 Jan 2024 18:18:35 +0100
+Subject: [PATCH] fix: allow building without git
+
+---
+ setup.py | 2 +-
+ 1 file changed, 1 insertion(+), 1 deletion(-)
+
+diff --git a/setup.py b/setup.py
+index e01c008..92eca62 100644
+--- a/setup.py
++++ b/setup.py
+@@ -71,7 +71,7 @@ def get_flash_version() -> str:
+             ["git", "describe", "--tags", "--always"],
+             cwd=flash_dir,
+         ).decode("ascii")[:-1]
+-    except subprocess.CalledProcessError:
++    except Exception:
+         version = flash_dir / "version.txt"
+         if version.is_file():
+             return version.read_text().strip()
+-- 
+2.43.0
+

pkgs/development/python-modules/xformers/default.nix

@@ -25,7 +25,8 @@
 #, flash-attn
 }:
 let
-  version = "0.03";
+  inherit (torch) cudaCapabilities cudaPackages cudaSupport;
+  version = "0.0.23.post1";
 in
 buildPythonPackage {
   pname = "xformers";
@@ -38,17 +39,34 @@ buildPythonPackage {
     owner = "facebookresearch";
     repo = "xformers";
     rev = "refs/tags/v${version}";
-    hash = "sha256-G8f7tny5B8SAQ6+2uOjhY7nD0uOT4sskIwtTdwivQXo=";
+    hash = "sha256-AJXow8MmX4GxtEE2jJJ/ZIBr+3i+uS4cA6vofb390rY=";
     fetchSubmodules = true;
   };
 
+  patches = [
+    ./0001-fix-allow-building-without-git.patch
+  ];
+
   preBuild = ''
     cat << EOF > ./xformers/version.py
     # noqa: C801
     __version__ = "${version}"
     EOF
+  '' + lib.optionalString cudaSupport ''
+    export CUDA_HOME=${cudaPackages.cuda_nvcc}
+    export TORCH_CUDA_ARCH_LIST="${lib.concatStringsSep ";" cudaCapabilities}"
   '';
 
+  buildInputs = lib.optionals cudaSupport (with cudaPackages; [
+    # flash-attn build
+    cuda_cudart # cuda_runtime_api.h
+    libcusparse.dev # cusparse.h
+    cuda_cccl.dev # nv/target
+    libcublas.dev # cublas_v2.h
+    libcusolver.dev # cusolverDn.h
+    libcurand.dev # curand_kernel.h
+  ]);
+
   nativeBuildInputs = [
     which
   ];

pkgs/top-level/python-packages.nix

@@ -16309,9 +16309,7 @@ self: super: with self; {
     inherit (pkgs) graphviz;
   };
 
-  xformers = callPackage ../development/python-modules/xformers {
-    openai-triton = self.openai-triton-cuda;
-  };
+  xformers = callPackage ../development/python-modules/xformers { };
 
   xgboost = callPackage ../development/python-modules/xgboost {
     inherit (pkgs) xgboost;