get_kernel: support specifying the backend (#268)

* `get_kernel`: support specifying the backend

So far, `get_kernel` has always auto-detected the backend from the Torch
build. However, this causes issues in fixed-device settings where we may
want to run some ops on the CPU and others on the GPU. This change
allows specifying `backend` as an argument to `get_kernel` and friends.

Currently, only the `cpu` backend and the backend that Torch was built
for are supported (since we cannot determine the version, etc. of a
backend that Torch was not built for). However, we decided to make it a
string-based argument as opposed to e.g. a `cpu` bool, in the case Torch
might do multi-backend builds in the future.

* Fix typing and incorrect `hip` backend name

* Test kernels package with Torch 2.9 and 2.10

* Do not attempt to lock Torch version

* Simplify `_select_backend`

* CI: try to fix

Author: danieldk

.github/workflows/test_kernels.yaml

@@ -24,7 +24,7 @@ jobs:
       max-parallel: 4
       matrix:
         python-version: ["3.10", "3.12"]
-        torch-version: ["2.7.0", "2.8.0"]
+        torch-version: ["2.9.0", "2.10.0"]
 
     env:
       UV_PYTHON_PREFERENCE: only-managed
@@ -38,14 +38,14 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Lock Torch version
-        working-directory: ./kernels
-        run: uv lock --upgrade-package "torch==${{ matrix.torch-version }}"
-
       - name: Install the project
         working-directory: ./kernels
         run: uv sync --all-extras --dev
 
+      - name: Install Torch version
+        working-directory: ./kernels
+        run: uv pip install "torch==${{ matrix.torch-version }}"
+
       - name: Install setuptools for Triton-based test
         working-directory: ./kernels
         run: uv pip install setuptools
@@ -76,7 +76,7 @@ jobs:
           HF_TOKEN: ${{ secrets.HF_STAGING_TOKEN }}
         run: |
           HUGGINGFACE_CO_STAGING=true uv run pytest --token -m "is_staging_test" tests/
-        if: matrix.python_version == '3.10' && matrix.torch-version == '2.7.0'
+        if: matrix.python_version == '3.10' && matrix.torch-version == '2.9.0'
 
       - name: Check README generation
         working-directory: ./kernels

kernels/src/kernels/cli/benchmark.py

@@ -20,7 +20,7 @@
 )
 
 from kernels.benchmark import Benchmark
-from kernels.utils import _get_hf_api, backend
+from kernels.utils import _get_hf_api, _backend
 
 MISSING_DEPS: list[str] = []
 
@@ -387,8 +387,8 @@ def collect_machine_info() -> MachineInfo:
 
     if TORCH_AVAILABLE:
         pytorch_version = torch.__version__
-        backend_name = backend()
-        if backend_name in {"cuda", "hip"}:
+        backend_name = _backend()
+        if backend_name in {"cuda", "rocm"}:
             gpu = torch.cuda.get_device_name(0)
             # ROCm uses the CUDA API but has torch.version.hip
             if hasattr(torch.version, "hip") and torch.version.hip:
@@ -465,10 +465,10 @@ def run_benchmark_class(
         kernel = get_kernel(repo_id, revision=revision)
 
     kernel_sha = get_kernel_sha_from_build_name(kernel)
-    backend_name = backend() if TORCH_AVAILABLE else "cpu"
+    backend_name = _backend() if TORCH_AVAILABLE else "cpu"
     # Map backend names to torch device names
     device_map = {
-        "hip": "cuda",
+        "rocm": "cuda",
         "metal": "mps",
         "cann": "npu",
     }

kernels/src/kernels/cli/versions.py

@@ -3,7 +3,7 @@
 from huggingface_hub import HfApi
 
 from kernels._versions import _get_available_versions
-from kernels.utils import _get_hf_api, build_variants
+from kernels.utils import _get_hf_api, _build_variants
 from kernels.variants import BUILD_VARIANT_REGEX
 
 
@@ -14,7 +14,7 @@ def print_kernel_versions(repo_id: str):
         # Do not mark compatible variants when Torch is not available.
         compatible_variants = set()
     else:
-        compatible_variants = set(build_variants())
+        compatible_variants = set(_build_variants(None))
 
     versions = _get_available_versions(repo_id).items()
     if not versions:

kernels/src/kernels/utils.py

@@ -71,13 +71,13 @@ def _get_privateuse_backend_name() -> str | None:
     return None
 
 
-def backend() -> str:
+def _backend() -> str:
     import torch
 
     if torch.version.cuda is not None:
         return "cuda"
     elif torch.version.hip is not None:
-        return "hip"
+        return "rocm"
     elif torch.backends.mps.is_available():
         return "metal"
     elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
@@ -88,21 +88,23 @@ def backend() -> str:
         return "cpu"
 
 
-def build_variant() -> str:
+def _build_variant(backend: str | None) -> str:
+    backend = _select_backend(backend)
+
     import torch
 
-    if torch.version.cuda is not None:
+    if backend == "cuda" and torch.version.cuda is not None:
         cuda_version = parse(torch.version.cuda)
         compute_framework = f"cu{cuda_version.major}{cuda_version.minor}"
-    elif torch.version.hip is not None:
+    elif backend == "rocm" and torch.version.hip is not None:
         rocm_version = parse(torch.version.hip.split("-")[0])
         compute_framework = f"rocm{rocm_version.major}{rocm_version.minor}"
-    elif torch.backends.mps.is_available():
+    elif backend == "metal":
         compute_framework = "metal"
-    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+    elif backend == "xpu" and torch.version.xpu is not None:
         version = torch.version.xpu
         compute_framework = f"xpu{version[0:4]}{version[5:6]}"
-    elif _get_privateuse_backend_name() == "npu":
+    elif backend == "cann":
         from torch_npu.utils.collect_env import get_cann_version  # type: ignore[import-not-found]
 
         cann_major, cann_minor = get_cann_version()[0], get_cann_version()[2]
@@ -125,36 +127,57 @@ def build_variant() -> str:
     return f"torch{torch_version.major}{torch_version.minor}-{cxxabi}-{compute_framework}-{cpu}-{os}"
 
 
-def build_variant_noarch() -> str:
-    import torch
+def _supported_backends() -> set[str]:
+    return {"cpu", _backend()}
 
-    if torch.version.cuda is not None:
+
+def _select_backend(backend: str | None) -> str:
+    if backend is None:
+        return _backend()
+
+    supported = _supported_backends()
+    if backend in supported:
+        return backend
+
+    raise ValueError(
+        f"Invalid backend '{backend}', system supported backends: {', '.join(sorted(supported))}"
+    )
+
+
+def _build_variant_noarch(backend: str | None) -> str:
+    backend = _select_backend(backend)
+
+    if backend == "cuda":
         return "torch-cuda"
-    elif torch.version.hip is not None:
+    elif backend == "rocm":
         return "torch-rocm"
-    elif torch.backends.mps.is_available():
+    elif backend == "metal":
         return "torch-metal"
-    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+    elif backend == "xpu":
         return "torch-xpu"
-    elif _get_privateuse_backend_name() == "npu":
+    elif backend == "cann":
         return "torch-npu"
     else:
         return "torch-cpu"
 
 
-def build_variant_universal() -> str:
+def _build_variant_universal() -> str:
     # Once we support other frameworks, detection goes here.
     return "torch-universal"
 
 
-def build_variants() -> list[str]:
+def _build_variants(backend: str | None) -> list[str]:
     """Return compatible build variants in preferred order."""
-    return [build_variant(), build_variant_noarch(), build_variant_universal()]
+    return [
+        _build_variant(backend),
+        _build_variant_noarch(backend),
+        _build_variant_universal(),
+    ]
 
 
 def _import_from_path(module_name: str, variant_path: Path) -> ModuleType:
     metadata = Metadata.load_from_variant(variant_path)
-    validate_dependencies(metadata.python_depends, backend())
+    validate_dependencies(metadata.python_depends, _backend())
 
     file_path = variant_path / "__init__.py"
     if not file_path.exists():
@@ -181,6 +204,7 @@ def install_kernel(
     repo_id: str,
     revision: str,
     local_files_only: bool = False,
+    backend: str | None = None,
     variant_locks: dict[str, VariantLock] | None = None,
     user_agent: str | dict | None = None,
 ) -> tuple[str, Path]:
@@ -196,6 +220,9 @@ def install_kernel(
             The specific revision (branch, tag, or commit) to download.
         local_files_only (`bool`, *optional*, defaults to `False`):
             Whether to only use local files and not download from the Hub.
+        backend (`str`, *optional*):
+            The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for.
+            The backend will be detected automatically if not provided.
         variant_locks (`dict[str, VariantLock]`, *optional*):
             Optional dictionary of variant locks for validation.
         user_agent (`Union[str, dict]`, *optional*):
@@ -205,7 +232,7 @@ def install_kernel(
         `tuple[str, Path]`: A tuple containing the package name and the path to the variant directory.
     """
     package_name = package_name_from_repo_id(repo_id)
-    allow_patterns = [f"build/{variant}/*" for variant in build_variants()]
+    allow_patterns = [f"build/{variant}/*" for variant in _build_variants(backend)]
     api = _get_hf_api(user_agent=user_agent)
     repo_path = Path(
         str(
@@ -220,7 +247,9 @@ def install_kernel(
     )
 
     try:
-        return _find_kernel_in_repo_path(repo_path, package_name, variant_locks)
+        return _find_kernel_in_repo_path(
+            repo_path, package_name, backend=backend, variant_locks=variant_locks
+        )
     except FileNotFoundError:
         raise FileNotFoundError(
             f"Cannot install kernel from repo {repo_id} (revision: {revision})"
@@ -230,9 +259,11 @@ def install_kernel(
 def _find_kernel_in_repo_path(
     repo_path: Path,
     package_name: str,
+    *,
+    backend: str | None = None,
     variant_locks: dict[str, VariantLock] | None = None,
 ) -> tuple[str, Path]:
-    variants = build_variants()
+    variants = _build_variants(backend)
     variant = None
     variant_path = None
     for candidate_variant in variants:
@@ -303,6 +334,7 @@ def get_kernel(
     repo_id: str,
     revision: str | None = None,
     version: int | str | None = None,
+    backend: str | None = None,
     user_agent: str | dict | None = None,
 ) -> ModuleType:
     """
@@ -319,6 +351,9 @@ def get_kernel(
         version (`int|str`, *optional*):
             The kernel version to download as an integer. The `str` variant is deprecated and will be
             removed in a future release. Cannot be used together with `revision`.
+        backend (`str`, *optional*):
+            The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for.
+            The backend will be detected automatically if not provided.
         user_agent (`Union[str, dict]`, *optional*):
             The `user_agent` info to pass to `snapshot_download()` for internal telemetry.
 
@@ -342,12 +377,16 @@ def get_kernel(
 
     revision = select_revision_or_version(repo_id, revision=revision, version=version)
     package_name, variant_path = install_kernel(
-        repo_id, revision=revision, user_agent=user_agent
+        repo_id, revision=revision, backend=backend, user_agent=user_agent
     )
     return _import_from_path(package_name, variant_path)
 
 
-def get_local_kernel(repo_path: Path, package_name: str) -> ModuleType:
+def get_local_kernel(
+    repo_path: Path,
+    package_name: str,
+    backend: str | None = None,
+) -> ModuleType:
     """
     Import a kernel from a local kernel repository path.
 
@@ -356,13 +395,16 @@ def get_local_kernel(repo_path: Path, package_name: str) -> ModuleType:
             The local path to the kernel repository.
         package_name (`str`):
             The name of the package to import from the repository.
+        backend (`str`, *optional*):
+            The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for.
+            The backend will be detected automatically if not provided.
 
     Returns:
         `ModuleType`: The imported kernel module.
     """
     # Presume we were given the top level path of the kernel repository.
     for base_path in [repo_path, repo_path / "build"]:
-        for v in build_variants():
+        for v in _build_variants(backend):
             variant_path = base_path / v
             if variant_path.exists():
                 return _import_from_path(package_name, variant_path)
@@ -377,7 +419,10 @@ def get_local_kernel(repo_path: Path, package_name: str) -> ModuleType:
 
 
 def has_kernel(
-    repo_id: str, revision: str | None = None, version: int | str | None = None
+    repo_id: str,
+    revision: str | None = None,
+    version: int | str | None = None,
+    backend: str | None = None,
 ) -> bool:
     """
     Check whether a kernel build exists for the current environment (Torch version and compute framework).
@@ -390,17 +435,19 @@ def has_kernel(
         version (`int|str`, *optional*):
             The kernel version to download as an integer. The `str` variant is deprecated and will be
             removed in a future release. Cannot be used together with `revision`.
+        backend (`str`, *optional*):
+            The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for.
+            The backend will be detected automatically if not provided.
 
     Returns:
         `bool`: `True` if a kernel is available for the current environment.
     """
     revision = select_revision_or_version(repo_id, revision=revision, version=version)
 
     package_name = package_name_from_repo_id(repo_id)
-    variant = build_variant()
 
     api = _get_hf_api()
-    for variant in build_variants():
+    for variant in _build_variants(backend):
         for init_file in ["__init__.py", f"{package_name}/__init__.py"]:
             if api.file_exists(
                 repo_id, revision=revision, filename=f"build/{variant}/{init_file}"
@@ -410,7 +457,12 @@ def has_kernel(
     return False
 
 
-def load_kernel(repo_id: str, *, lockfile: Path | None) -> ModuleType:
+def load_kernel(
+    repo_id: str,
+    *,
+    lockfile: Path | None,
+    backend: str | None = None,
+) -> ModuleType:
     """
     Get a pre-downloaded, locked kernel.
 
@@ -421,6 +473,9 @@ def load_kernel(repo_id: str, *, lockfile: Path | None) -> ModuleType:
             The Hub repository containing the kernel.
         lockfile (`Path`, *optional*):
             Path to the lockfile. If not provided, the lockfile will be loaded from the caller's package metadata.
+        backend (`str`, *optional*):
+            The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for.
+            The backend will be detected automatically if not provided.
 
     Returns:
         `ModuleType`: The imported kernel module.
@@ -439,7 +494,7 @@ def load_kernel(repo_id: str, *, lockfile: Path | None) -> ModuleType:
     package_name = package_name_from_repo_id(repo_id)
 
     api = _get_hf_api()
-    allow_patterns = [f"build/{variant}/*" for variant in build_variants()]
+    allow_patterns = [f"build/{variant}/*" for variant in _build_variants(backend)]
     repo_path = Path(
         str(
             api.snapshot_download(
@@ -454,7 +509,7 @@ def load_kernel(repo_id: str, *, lockfile: Path | None) -> ModuleType:
 
     try:
         package_name, variant_path = _find_kernel_in_repo_path(
-            repo_path, package_name, variant_locks=None
+            repo_path, package_name, backend=backend, variant_locks=None
         )
         return _import_from_path(package_name, variant_path)
     except FileNotFoundError:
@@ -605,7 +660,7 @@ def _get_hf_api(user_agent: str | dict | None = None) -> HfApi:
 
         # System info
         python = ".".join(platform.python_version_tuple()[:2])
-        user_agent_str += f"; kernels/{__version__}; python/{python}; torch/{torch.__version__}; build_variant/{build_variant()}; file_type/kernel"
+        user_agent_str += f"; kernels/{__version__}; python/{python}; torch/{torch.__version__}; build_variant/{_build_variant(None)}; file_type/kernel"
 
         # Add glibc version if available
         glibc = glibc_version()