diff --git a/flake.lock b/flake.lock index 719cdeea0..82f649fdd 100644 --- a/flake.lock +++ b/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" } diff --git a/flake.nix b/flake.nix index 5058667a8..ac4331c24 100644 --- 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/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 = { diff --git a/nix/server.nix b/nix/server.nix index 0640fe3a1..bc572ccda 100644 --- a/nix/server.nix +++ b/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 diff --git a/server/pyproject.toml b/server/pyproject.toml index 07ea1048f..1c26418f4 100644 --- a/server/pyproject.toml +++ b/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] diff --git a/server/text_generation_server/utils/kernels.py b/server/text_generation_server/utils/kernels.py index 42745c716..b6f43259b 100644 --- a/server/text_generation_server/utils/kernels.py +++ b/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"]