Set version to 0.1.3 (#17)

* Only import torch when needed

This avoids the (costly) torch load when e.g. the setuptools hooks
are running in downstream packages.

* Lock Python/Torch versions

Also update to Torch 2.5.1/2.6.0.

* Set the minimum Python version to 3.9

* Change step description

.github/workflows/test.yml

@@ -23,8 +23,8 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python: ["3.10", "3.12"]
-        torch: ["2.4.0", "2.5.0"]
+        python-version: ["3.10", "3.12"]
+        torch-version: ["2.5.1", "2.6.0"]
 
     steps:
       - name: Checkout code
@@ -35,6 +35,9 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Lock Torch version
+        run: uv lock --upgrade-package "torch==${{ matrix.torch-version }}"
+
       - name: Install the project
         run: uv sync --all-extras --dev
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hf-kernels"
-version = "0.1.2"
+version = "0.1.3"
 description = "Download cuda kernels"
 authors = [
   { name = "OlivierDehaene", email = "olivier@huggingface.co" },
@@ -9,7 +9,7 @@ authors = [
   { name = "Nicolas Patry", email = "nicolas@huggingface.co" },
 ]
 readme = "README.md"
-requires-python = ">= 3.8"
+requires-python = ">= 3.9"
 dependencies = [
   "huggingface-hub>=0.26.3",
   "packaging>=24.2",

src/hf_kernels/cli.py

@@ -6,14 +6,14 @@ from pathlib import Path
 
 from hf_kernels.compat import tomllib
 from hf_kernels.lockfile import KernelLock, get_kernel_locks
-from hf_kernels.utils import install_kernel
+from hf_kernels.utils import install_kernel, install_kernel_all_variants
 
 
 def main():
     parser = argparse.ArgumentParser(
         prog="hf-kernel", description="Manage compute kernels"
     )
-    subparsers = parser.add_subparsers()
+    subparsers = parser.add_subparsers(required=True)
 
     download_parser = subparsers.add_parser("download", help="Download locked kernels")
     download_parser.add_argument(
@@ -21,6 +21,11 @@ def main():
         type=Path,
         help="The project directory",
     )
+    download_parser.add_argument(
+        "--all-variants",
+        action="store_true",
+        help="Download all build variants of the kernel",
+    )
     download_parser.set_defaults(func=download_kernels)
 
     lock_parser = subparsers.add_parser("lock", help="Lock kernel revisions")
@@ -51,7 +56,10 @@ def download_kernels(args):
             f"Downloading `{kernel_lock.repo_id}` at with SHA: {kernel_lock.sha}",
             file=sys.stderr,
         )
-        install_kernel(kernel_lock.repo_id, kernel_lock.sha)
+        if args.all_variants:
+            install_kernel_all_variants(kernel_lock.repo_id, kernel_lock.sha)
+        else:
+            install_kernel(kernel_lock.repo_id, kernel_lock.sha)
 
 
 def lock_kernels(args):

src/hf_kernels/utils.py

@@ -9,7 +9,6 @@ from importlib.metadata import Distribution
 from types import ModuleType
 from typing import List, Optional
 
-import torch
 from huggingface_hub import hf_hub_download, snapshot_download
 from packaging.version import parse
 
@@ -18,6 +17,8 @@ from hf_kernels.lockfile import KernelLock
 
 
 def build_variant():
+    import torch
+
     torch_version = parse(torch.__version__)
     cuda_version = parse(torch.version.cuda)
     cxxabi = "cxx11" if torch.compiled_with_cxx11_abi() else "cxx98"
@@ -48,6 +49,17 @@ def install_kernel(repo_id: str, revision: str, local_files_only: bool = False):
     return package_name, f"{repo_path}/build/{build_variant()}"
 
 
+def install_kernel_all_variants(
+    repo_id: str, revision: str, local_files_only: bool = False
+):
+    snapshot_download(
+        repo_id,
+        allow_patterns="build/*",
+        revision=revision,
+        local_files_only=local_files_only,
+    )
+
+
 def get_metadata(repo_id: str, revision: str, local_files_only: bool = False):
     with open(
         hf_hub_download(