nix-builder: wire up kernel dependencies

docs/source/builder/writing-kernels.md

@@ -2,7 +2,7 @@
 
 > [!TIP]
 > Refer to the [kernel requirements' page](../kernel-requirements.md) to
-> get an idea of what is expected from the kernel structure and content. 
+> get an idea of what is expected from the kernel structure and content.
 
 ## Introduction
 
@@ -320,7 +320,7 @@ The following options can be set for a kernel:
   supported backends are `cpu`, `cuda`, `metal`, `rocm`, and `xpu`.
   **The `cpu` backend is currently experimental and might still change.**
 - `depends` (required): a list of dependencies. The supported dependencies
-  are listed in [`deps.nix`](https://github.com/huggingface/kernels/blob/main/builder/lib/deps.nix).
+  are listed in [`cpp-deps.nix`](https://github.com/huggingface/kernels/blob/main/builder/lib/cpp-deps.nix).
 - `src` (required): a list of source files and headers.
 - `include` (optional): include directories relative to the project root.
   Default: `[]`.

examples/kernels/flake.nix

@@ -207,6 +207,11 @@
           path = ./gemm-triton-autotune;
           drv = sys: out: out.packages.${sys}.redistributable.torch-cuda;
         }
+        {
+          name = "kernel-deps-kernel";
+          path = ./kernel-deps;
+          drv = sys: out: out.packages.${sys}.redistributable.torch-cuda;
+        }
       ];
 
       # ROCm kernels to build in CI.

examples/kernels/kernel-deps/CARD.md

@@ -0,0 +1,66 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+# If the org / user isn't a trusted publisher, pass `trust_remote_code=True` to the
+# `get_kernel` call. You can find whether this kernel is from a trusted publisher
+# by going to the kernel's Hub page and finding the "Trusted publisher" status at
+# the top of the page.
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Upstream
+
+The original source code for this kernel comes from {{ upstream }}.
+{% endif %}
+{% if source %}
+
+## Source
+
+The kernel-builder formatted source for this kernel is available at {{ source }}.
+{% endif %}
+

examples/kernels/kernel-deps/build.toml

@@ -0,0 +1,24 @@
+[general]
+name = "kernel-deps"
+version = 1
+edition = 5
+license = "apache-2.0"
+backends = [
+    "cpu",
+    "cuda",
+    "metal",
+    "rocm",
+    "xpu",
+]
+
+kernel-depends = [
+    { repo-id = "kernels-community/activation", version = 1 },
+    { repo-id = "kernels-community/einops", version = 1 },
+]
+
+[general.hub]
+repo-id = "kernels-test/kernel-deps"
+
+[torch-noarch]
+
+[kernel]

examples/kernels/kernel-deps/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for kernels tests";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/kernel-deps/kernels.lock

@@ -0,0 +1,22 @@
+[
+  {
+    "dependency": {
+      "repo-id": "kernels-community/activation",
+      "version": 1
+    },
+    "lock": {
+      "commit": "36f4b285ab9125127d1b7364a786089ef10bcf8f",
+      "hash": "sha256-S2bSgxoK5JbuyBnzZSXNcV7NmMGtccuqfsjPgbwki7o="
+    }
+  },
+  {
+    "dependency": {
+      "repo-id": "kernels-community/einops",
+      "version": 1
+    },
+    "lock": {
+      "commit": "94986c119039ea0cb65592c210c57255de1089eb",
+      "hash": "sha256-qeQZSajDKBw4gfjZEE6wreqvdSKGfZmkwi/KnkjPj5g="
+    }
+  }
+]

examples/kernels/kernel-deps/tests/test_kernel_deps.py

@@ -0,0 +1,18 @@
+import torch
+import kernels
+
+kernel_deps = kernels.get_kernel("kernels-test/kernel-deps", version=1)
+
+
+def test_silu():
+    torch.manual_seed(42)
+    x = torch.randn(4, 64, device="cuda")
+    torch.testing.assert_close(kernel_deps.silu(x), torch.nn.functional.silu(x))
+
+
+def test_swap_last_dimensions():
+    torch.manual_seed(42)
+    x = torch.randn(3, 4, 5, 6)
+    y = kernel_deps.swap_last_dimensions(x)
+    assert y.shape == (3, 4, 6, 5)
+    torch.testing.assert_close(y, x.transpose(-1, -2))

examples/kernels/kernel-deps/torch-ext/kernel_deps/__init__.py

@@ -0,0 +1,33 @@
+import kernels
+import torch
+
+
+activation = kernels.get_kernel_dep("kernels-community/activation")
+einops = kernels.get_kernel_dep("kernels-community/einops")
+
+
+def silu(x: torch.Tensor) -> torch.Tensor:
+    out = torch.empty_like(x)
+    activation.silu(out, x)
+    return out
+
+
+def swap_last_dimensions(x: torch.Tensor) -> torch.Tensor:
+    """
+    Swap the last two dimensions of a tensor.
+
+    Shapes:
+        x: (..., M, N) with at least two dimensions
+        return: (..., N, M)
+
+    Raises:
+        ValueError: if the tensor has fewer than two dimensions.
+    """
+    if x.ndim < 2:
+        raise ValueError(
+            f"Expected a tensor with at least two dimensions, got {x.ndim}"
+        )
+    return einops.rearrange(x, "... a b -> ... b a")
+
+
+__all__ = ["swap_last_dimensions"]

nix-builder/lib/build.nix

@@ -136,6 +136,7 @@ rec {
           _: kernel: builtins.length (kernel.cuda-capabilities or supportedCudaCapabilities)
         ) kernelConfig.toml.kernel
       );
+      kernelDeps = pkgs.fetchKernelDeps src;
       pythonDeps = (kernelConfig.toml.general.python-depends or [ ]);
       backendPythonDeps =
         lib.attrByPath [ buildConfig.backend "python-depends" ] [ ]
@@ -150,6 +151,7 @@ rec {
           src
           rev
           doGetKernelCheck
+          kernelDeps
           pythonDeps
           backendPythonDeps
           kernelProvenance
@@ -166,6 +168,7 @@ rec {
           src
           stripRPath
           rev
+          kernelDeps
           pythonDeps
           backendPythonDeps
           kernelProvenance
@@ -184,6 +187,7 @@ rec {
           src
           stripRPath
           rev
+          kernelDeps
           pythonDeps
           backendPythonDeps
           kernelProvenance

nix-builder/lib/cpp-deps.nix

@@ -0,0 +1,47 @@
+{
+  lib,
+  pkgs,
+  stdenv,
+  torch,
+}:
+
+let
+  deps = {
+    "cutlass_2_10" = [
+      pkgs.cutlass_2_10
+    ];
+    "cutlass_3_5" = [
+      pkgs.cutlass_3_5
+    ];
+    "cutlass_3_6" = [
+      pkgs.cutlass_3_6
+    ];
+    "cutlass_3_8" = [
+      pkgs.cutlass_3_8
+    ];
+    "cutlass_3_9" = [
+      pkgs.cutlass_3_9
+    ];
+    "cutlass_4_0" = [
+      pkgs.cutlass_4_0
+    ];
+    "cutlass_4_4" = [
+      pkgs.cutlass_4_4
+    ];
+    "cutlass_4_5" = [
+      pkgs.cutlass_4_5
+    ];
+    "torch" = [
+      torch
+    ];
+    "sycl_tla" = [
+      (torch.xpuPackages.sycl-tla.override { inherit stdenv; })
+    ];
+    "metal-cpp" = [
+      (pkgs.metal-cpp.override { inherit stdenv; }).dev
+    ];
+  };
+
+  getCppDep = dep: deps.${dep} or (throw "Unknown dependency: ${dep}");
+in
+deps: lib.flatten (map getCppDep deps)

nix-builder/lib/extension/default.nix

@@ -91,13 +91,12 @@ in
 
   mkTorchNoArchExtension = callPackage ./torch/no-arch.nix { inherit torch; };
 
-  inherit
-    (import ../deps.nix {
+  resolveCppDeps = (
+    import ../cpp-deps.nix {
       inherit lib pkgs torch;
       stdenv = effectiveStdenv;
-    })
-    resolveCppDeps
-    ;
+    }
+  );
 
   stdenv = effectiveStdenv;
 }

nix-builder/lib/extension/torch/arch.nix

@@ -55,6 +55,10 @@
 
   nvccThreads,
 
+  # Dependencies on other kernels. Path to a JSON file that maps
+  # kernel dependencies to output paths.
+  kernelDeps,
+
   # A stringly-typed list of Python dependencies. Ideally we'd take a
   # list of derivations, but we also need to write the dependencies to
   # the output.
@@ -86,7 +90,7 @@ assert (buildConfig.metal or false) -> stdenv.hostPlatform.isDarwin;
 
 let
   inherit
-    (import ../../deps.nix {
+    (import ../../python-deps.nix {
       inherit
         lib
         pkgs
@@ -125,6 +129,7 @@ stdenv.mkDerivation (prevAttrs: {
   inherit
     doAbiCheck
     doKernelBuildCheck
+    kernelDeps
     moduleName
     nvccThreads
     ;

nix-builder/lib/extension/torch/no-arch.nix

@@ -35,6 +35,10 @@
 
   src,
 
+  # Dependencies on other kernels. Path to a JSON file that maps
+  # kernel dependencies to output paths.
+  kernelDeps,
+
   # A stringly-typed list of Python dependencies. Ideally we'd take a
   # list of derivations, but we also need to write the dependencies to
   # the output.
@@ -55,7 +59,7 @@ assert (buildConfig.metal or false) -> stdenv.hostPlatform.isDarwin;
 
 let
   inherit
-    (import ../../deps.nix {
+    (import ../../python-deps.nix {
       inherit
         lib
         pkgs
@@ -78,7 +82,7 @@ in
 stdenv.mkDerivation (prevAttrs: {
   name = "${kernelName}-torch-ext";
 
-  inherit doKernelBuildCheck moduleName;
+  inherit doKernelBuildCheck kernelDeps moduleName;
 
   src = pkgs.runCommand "source" { } ''
     mkdir -p $out

nix-builder/lib/extension/tvm-ffi/arch.nix

@@ -55,6 +55,10 @@
 
   nvccThreads,
 
+  # Dependencies on other kernels. Path to a JSON file that maps
+  # kernel dependencies to output paths.
+  kernelDeps,
+
   # A stringly-typed list of Python dependencies. Ideally we'd take a
   # list of derivations, but we also need to write the dependencies to
   # the output.
@@ -83,7 +87,7 @@ assert (buildConfig.metal or false) -> stdenv.hostPlatform.isDarwin;
 
 let
   inherit
-    (import ../../deps.nix {
+    (import ../../python-deps.nix {
       inherit
         lib
         pkgs
@@ -125,6 +129,7 @@ stdenv.mkDerivation (prevAttrs: {
   inherit
     doAbiCheck
     doKernelBuildCheck
+    kernelDeps
     moduleName
     nvccThreads
     ;