Use rotary kernel from the Hub (#3041)

huggingface · Feb 21, 2025 · 97c5f7e · 97c5f7e
1 parent 1cae319
commit 97c5f7e
Show file tree

Hide file tree

Showing 8 changed files with 195 additions and 109 deletions.
diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -5,7 +5,7 @@
       inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
     };
     nix-filter.url = "github:numtide/nix-filter";
-    tgi-nix.url = "github:huggingface/text-generation-inference-nix/flashinfer-0.2.0.post2";
+    tgi-nix.url = "github:huggingface/text-generation-inference-nix/hub-rotary";
     nixpkgs.follows = "tgi-nix/nixpkgs";
     flake-utils.url = "github:numtide/flake-utils";
     rust-overlay = {

diff --git a/nix/server.nix b/nix/server.nix
@@ -11,7 +11,6 @@
   flashinfer,
   flash-attn,
   flash-attn-layer-norm,
-  flash-attn-rotary,
   flash-attn-v1,
   grpc-interceptor,
   grpcio-reflection,
@@ -36,6 +35,7 @@
   pydantic,
   quantization,
   quantization-eetq,
+  rotary,
   safetensors,
   tokenizers,
   torch,
@@ -87,7 +87,6 @@ buildPythonPackage {
     flashinfer
     flash-attn
     flash-attn-layer-norm
-    flash-attn-rotary
     grpc-interceptor
     grpcio-reflection
     grpcio-status
@@ -111,6 +110,7 @@ buildPythonPackage {
     pydantic
     quantization
     quantization-eetq
+    rotary
     safetensors
     sentencepiece
     tokenizers