1 files changed, 73 insertions, 0 deletions
diff --git a/cuda/flashattn/flake.nix b/cuda/flashattn/flake.nix
new file mode 100644
index 0000000..978bff2
--- /dev/null
+++ b/cuda/flashattn/flake.nix
@@ -0,0 +1,73 @@
+{
+  description = "Torch cuda flake using nix-community cachix";
+
+  nixConfig = {
+    extra-substituters = [
+      "https://nix-community.cachix.org"
+      "https://nix-ai-stuff.cachix.org"
+      "https://ai.cachix.org"
+      "https://cuda-maintainers.cachix.org"
+    ];
+    extra-trusted-public-keys = [
+      "nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs="
+      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
+      "ai.cachix.org-1:N9dzRK+alWwoKXQlnn0H6aUx0lU/mspIoz8hMvGvbbc="
+      "nix-ai-stuff.cachix.org-1:WlUGeVCs26w9xF0/rjyg32PujDqbVMlSHufpj1fqix8="
+    ];
+  };
+  inputs = {
+    nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
+    nix-ai-stuff = {
+      url = "github:BatteredBunny/nix-ai-stuff";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+  };
+
+  outputs = {
+    self,
+    nixpkgs,
+    nix-ai-stuff,
+    ...
+  }: let
+    pkgs = import nixpkgs {
+      system = "x86_64-linux";
+      config.allowUnfree = true;
+      config.cudaSupport = true;
+    };
+  in {
+    devShell.x86_64-linux = with pkgs;
+      mkShell rec {
+        venvDir = "./.venv";
+        buildInputs = [
+          (pkgs.python3.withPackages (
+            ps:
+              with ps; [
+                torch
+                accelerate
+                transformers
+                typing-extensions
+                psutil
+                ninja
+                einops
+                packaging
+                sentence-transformers
+                nix-ai-stuff.packages.${pkgs.system}.flash-attn
+              ]
+          ))
+          pkgs.virtualenv
+          pkgs.python3Packages.venvShellHook
+        ];
+        lulz = [
+          pkgs.python3Packages.sentence-transformers
+          nix-ai-stuff.packages.${pkgs.system}.flash-attn
+        ];
+        HENLO = "${pkgs.lib.makeLibraryPath lulz}";
+        postVenvCreation = ''
+          unset SOURCE_DATE_EPOCH
+        '';
+        shellHook = ''
+          fish
+        '';
+      };
+  };
+}