nix: add flash-attn-v1 to the server environment

flake.lock

@@ -978,16 +978,16 @@
         "nixpkgs": "nixpkgs_6"
       },
       "locked": {
-        "lastModified": 1726743157,
-        "narHash": "sha256-7OczwJsA47o+aUftMwkoh8R31DlNSl2FgRjqE8zAggk=",
-        "owner": "danieldk",
-        "repo": "tgi-nix",
-        "rev": "bcc9fd01cf81bc42cebb999a736a377adfa8942f",
+        "lastModified": 1727353315,
+        "narHash": "sha256-yZovq/6P8Z199r7e+NbTXyCqRgK6grRkLxYHWHnHckI=",
+        "owner": "huggingface",
+        "repo": "text-generation-inference-nix",
+        "rev": "1d42c4125ebafb87707118168995675cc5050b9d",
         "type": "github"
       },
       "original": {
-        "owner": "danieldk",
-        "repo": "tgi-nix",
+        "owner": "huggingface",
+        "repo": "text-generation-inference-nix",
         "type": "github"
       }
     }

flake.nix

@@ -5,7 +5,7 @@
       inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
     };
     nix-filter.url = "github:numtide/nix-filter";
-    tgi-nix.url = "github:danieldk/tgi-nix";
+    tgi-nix.url = "github:huggingface/text-generation-inference-nix";
     nixpkgs.follows = "tgi-nix/nixpkgs";
     flake-utils.url = "github:numtide/flake-utils";
     rust-overlay = {
@@ -132,49 +132,12 @@
                 pre-commit
                 ruff
               ]);
-
           };
 
-          impure = mkShell {
-            buildInputs =
-              [
-                openssl.dev
-                pkg-config
-                (rust-bin.stable.latest.default.override {
-                  extensions = [
-                    "rust-analyzer"
-                    "rust-src"
-                  ];
-                })
-                protobuf
-              ]
-              ++ (with python3.pkgs; [
-                venvShellHook
-                docker
-                pip
-                ipdb
-                click
-                pyright
-                pytest
-                pytest-asyncio
-                redocly
-                ruff
-                syrupy
-              ]);
+          impure = callPackage ./nix/impure-shell.nix { inherit server; };
 
-            inputsFrom = [ server ];
-
-            venvDir = "./.venv";
-
-            postVenvCreation = ''
-              unset SOURCE_DATE_EPOCH
-              ( cd server ; python -m pip install --no-dependencies -e . )
-              ( cd clients/python ; python -m pip install --no-dependencies -e . )
-            '';
-            postShellHook = ''
-              unset SOURCE_DATE_EPOCH
-              export PATH=$PATH:~/.cargo/bin
-            '';
+          impure-flash-attn-v1 = callPackage ./nix/impure-shell.nix {
+            server = server.override { flash-attn = python3.pkgs.flash-attn-v1; };
           };
         };
 

nix/impure-shell.nix

@@ -0,0 +1,54 @@
+{
+  mkShell,
+  openssl,
+  pkg-config,
+  protobuf,
+  python3,
+  pyright,
+  redocly,
+  ruff,
+  rust-bin,
+  server,
+}:
+
+mkShell {
+  buildInputs =
+    [
+      openssl.dev
+      pkg-config
+      (rust-bin.stable.latest.default.override {
+        extensions = [
+          "rust-analyzer"
+          "rust-src"
+        ];
+      })
+      protobuf
+      pyright
+      redocly
+      ruff
+    ]
+    ++ (with python3.pkgs; [
+      venvShellHook
+      docker
+      pip
+      ipdb
+      click
+      pytest
+      pytest-asyncio
+      syrupy
+    ]);
+
+  inputsFrom = [ server ];
+
+  venvDir = "./.venv";
+
+  postVenvCreation = ''
+    unset SOURCE_DATE_EPOCH
+    ( cd server ; python -m pip install --no-dependencies -e . )
+    ( cd clients/python ; python -m pip install --no-dependencies -e . )
+  '';
+  postShellHook = ''
+    unset SOURCE_DATE_EPOCH
+    export PATH=$PATH:~/.cargo/bin
+  '';
+}

nix/server.nix

@@ -13,6 +13,7 @@
   flash-attn,
   flash-attn-layer-norm,
   flash-attn-rotary,
+  flash-attn-v1,
   grpc-interceptor,
   grpcio-reflection,
   grpcio-status,