Update to kernels 0.2.0.

flake.lock

@@ -853,11 +853,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1741141853,
-        "narHash": "sha256-FauVtC+FbOgkKpGVuQTNxSqrvgbmVc7hFkjn/DacwMo=",
+        "lastModified": 1741833135,
+        "narHash": "sha256-HUtFcF4NLwvu7CAowWgqCHXVkNj0EOc/W6Ism4biV6I=",
         "owner": "oxalica",
         "repo": "rust-overlay",
-        "rev": "02edad1f19d6dec824e0812e4cdc0aa7930ff8ae",
+        "rev": "f3cd1e0feb994188fe3ad9a5c3ab021ed433b8c8",
         "type": "github"
       },
       "original": {
@@ -978,16 +978,16 @@
         "nixpkgs": "nixpkgs_6"
       },
       "locked": {
-        "lastModified": 1740049068,
-        "narHash": "sha256-heYzYOt+TSnRKHIV24s74yEjLkTbBfjNCWHdQEX++eI=",
+        "lastModified": 1741617161,
+        "narHash": "sha256-cwKYAsIVSLtoLbG48+oi3NkSrvuZRLYs8lkJmpDsTw0=",
         "owner": "huggingface",
         "repo": "text-generation-inference-nix",
-        "rev": "143e8451efa22b120f97e6698508e9a0aed82769",
+        "rev": "5946021ec6cb6aae18158a9dc27f893cfbab2925",
         "type": "github"
       },
       "original": {
         "owner": "huggingface",
-        "ref": "hub-rotary",
+        "ref": "kernels-0.2.0",
         "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: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 = {

nix/server.nix

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

server/pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
     "grpcio>=1.67.0",
     "grpcio-reflection>=1.67.0",
     "grpcio-status>=1.67.0",
-    "hf-kernels>=0.1.5",
+    "kernels>=0.2.0",
     "hf-transfer>=0.1.8",
     "loguru>=0.7.3",
     "numpy>=1.26,<3",
@@ -36,7 +36,7 @@ dependencies = [
 ]
 
 [build-system]
-requires = ["hf-kernels>=0.1.2", "setuptools"]
+requires = ["kernels>=0.2.0", "setuptools"]
 build-backend = "setuptools.build_meta"
 
 [tool.kernels.dependencies]

server/text_generation_server/utils/kernels.py

@@ -1,7 +1,7 @@
 import importlib
 
 from loguru import logger
-from hf_kernels import load_kernel as hf_load_kernel
+from kernels import load_kernel as hf_load_kernel
 
 from text_generation_server.utils.log import log_once
 
@@ -16,7 +16,7 @@ def load_kernel(*, module: str, repo_id: str):
         log_once(logger.info, f"Using local module for `{module}`")
         return m
     except ModuleNotFoundError:
-        return hf_load_kernel(repo_id=repo_id)
+        return hf_load_kernel(repo_id=repo_id, module=module)
 
 
 __all__ = ["load_kernel"]