Update to kernels 0.2.1 (#3084)

* Update to `kernels` 0.2.1

The package was renamed from `hf-kernels` to `kernels`. The new version
also updates the lockfile format.

* Download kernels in `install-cuda` target

Author: danieldk

.github/workflows/tests.yaml

@@ -46,7 +46,7 @@ jobs:
       - name: Download locked kernels
         run: |
           source ./.venv/bin/activate
-          hf-kernels download server
+          kernels download server
       - name: Run server tests
         run: |
           source ./.venv/bin/activate

Dockerfile

@@ -183,12 +183,12 @@ COPY server server
 COPY server/Makefile server/Makefile
 ENV HF_KERNELS_CACHE=/kernels
 RUN cd server && \
-	uv sync --frozen --extra gen --extra attention --extra bnb --extra accelerate --extra compressed-tensors --extra marlin --extra moe --extra quantize --extra peft --extra outlines --no-install-project --active && \
+	uv sync --frozen --extra gen --extra bnb --extra accelerate --extra compressed-tensors --extra quantize --extra peft --extra outlines --no-install-project --active && \
     make gen-server-raw && \
-    hf-kernels download .
+    kernels download .
 
 RUN cd server && \
-    uv sync --frozen --extra gen --extra attention --extra bnb --extra accelerate --extra compressed-tensors --extra marlin --extra moe --extra quantize --extra peft --extra outlines --active --python=${PYTHON_VERSION} && \
+    uv sync --frozen --extra gen --extra bnb --extra accelerate --extra compressed-tensors --extra quantize --extra peft --extra outlines --active --python=${PYTHON_VERSION} && \
     uv pip install nvidia-nccl-cu12==2.25.1 && \
     pwd && \
     text-generation-server --help

flake.lock

@@ -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/hub-rotary";
+    tgi-nix.url = "github:huggingface/text-generation-inference-nix/kernels-0.2.0";
     nixpkgs.follows = "tgi-nix/nixpkgs";
     flake-utils.url = "github:numtide/flake-utils";
     rust-overlay = {

flake.nix

@@ -16,8 +16,8 @@
   grpcio-reflection,
   grpcio-status,
   grpcio-tools,
-  hf-kernels,
   hf-transfer,
+  kernels,
   loguru,
   mamba-ssm,
   moe,
@@ -91,8 +91,8 @@ buildPythonPackage {
     grpcio-reflection
     grpcio-status
     grpcio-tools
-    hf-kernels
     hf-transfer
+    kernels
     loguru
     mamba-ssm
     moe

nix/server.nix

@@ -39,6 +39,7 @@ install: install-cuda
 install-cuda: install-server install-flash-attention-v2-cuda install-flash-attention
 	uv pip install -e ".[attention,bnb,marlin,moe]"
 	uv pip install nvidia-nccl-cu12==2.22.3
+	kernels download .
 
 install-rocm: install-server install-flash-attention-v2-rocm  install-vllm-rocm