Merge branch 'main' into fix-metal-mps-lifecycle

robtaylor · web-flow · commit 0913ee421da7 · 2026-03-09T14:49:04.000Z
diff --git a/build2cmake/Cargo.lock b/build2cmake/Cargo.lock
diff --git a/build2cmake/Cargo.toml b/build2cmake/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "build2cmake"
-version = "0.12.2-dev0"
+version = "0.13.0-dev0"
 edition = "2021"
 description = "Generate CMake files for kernel-builder projects"
 homepage = "https://github.com/huggingface/kernel-builder"
diff --git a/builder/lib/cache.nix b/builder/lib/cache.nix
@@ -22,6 +22,7 @@
           ++ allOutputs build2cmake
           ++ allOutputs kernel-abi-check
           ++ allOutputs python3Packages.kernels
+          ++ allOutputs python3Packages.tvm-ffi
           ++ lib.optionals stdenv.hostPlatform.isLinux (allOutputs stdenvGlibc_2_27)
         );
       buildSetLinkFarm = buildSet: pkgs.linkFarm buildSet.torch.variant (buildSetOutputs buildSet);
diff --git a/builder/lib/deps.nix b/builder/lib/deps.nix
@@ -35,7 +35,7 @@ let
 
   pythonDeps =
     let
-      depsJson = builtins.fromJSON (builtins.readFile ../build2cmake/src/python_dependencies.json);
+      depsJson = builtins.fromJSON (builtins.readFile ../../build2cmake/src/python_dependencies.json);
       # Map the Nix package names to actual Nix packages.
       updatePackage = _name: dep: dep // { nix = map (pkg: pkgs.python3.pkgs.${pkg}) dep.nix; };
       updateBackend = _backend: backendDeps: lib.mapAttrs updatePackage backendDeps;
diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md
@@ -35,6 +35,16 @@ as the running example. After reading this page, you may also want to have
 a look at the more realistic [ReLU kernel with backprop and `torch.compile`](https://github.com/huggingface/kernels/tree/main/builder/examples/relu-backprop-compile)
 support.
 
+## Setting up environment
+
+In the [`terraform`](../../../terraform/) directory, we provide an
+example of programatically spinning up an EC2 instance that is ready
+with everything needed for you to start developing and building
+kernels.
+
+If you use a different provider, the Terraform bridges should be
+similar and straightforward to modify.
+
 ## Kernel project layout
 
 Kernel projects follow this general directory layout:
diff --git a/flake.nix b/flake.nix
@@ -169,6 +169,7 @@
               tabulate
               tomlkit
               torch
+              tvm-ffi
               types-pyyaml
               types-requests
               types-tabulate
diff --git a/kernel-abi-check/bindings/python/Cargo.lock b/kernel-abi-check/bindings/python/Cargo.lock
diff --git a/kernel-abi-check/bindings/python/Cargo.toml b/kernel-abi-check/bindings/python/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "kernel-abi-check-python"
-version = "0.12.2-dev0"
+version = "0.13.0-dev0"
 edition = "2024"
 description = "Check the ABI of Hub Kernels"
 homepage = "https://github.com/huggingface/kernel-builder"
diff --git a/kernel-abi-check/kernel-abi-check/Cargo.lock b/kernel-abi-check/kernel-abi-check/Cargo.lock
diff --git a/kernel-abi-check/kernel-abi-check/Cargo.toml b/kernel-abi-check/kernel-abi-check/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "kernel-abi-check"
-version = "0.12.2-dev0"
+version = "0.13.0-dev0"
 edition = "2021"
 description = "Check the ABI of Hub Kernels"
 homepage = "https://github.com/huggingface/kernel-builder"
diff --git a/kernels/pyproject.toml b/kernels/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "kernels"
-version = "0.12.2.dev0"
+version = "0.13.0.dev0"
 description = "Download compute kernels"
 authors = [
   { name = "Daniel de Kok", email = "daniel@huggingface.co" },
diff --git a/kernels/uv.lock b/kernels/uv.lock
diff --git a/nix/overlay.nix b/nix/overlay.nix
@@ -109,6 +109,9 @@ in
           torchVersion = "2.9";
           xpuPackages = final.xpuPackages_2025_2;
         };
+
+        tvm-ffi = callPackage ./pkgs/python-modules/tvm-ffi {
+        };
       }
     )
     (import ./pkgs/python-modules/hooks)
diff --git a/nix/pkgs/python-modules/tvm-ffi/default.nix b/nix/pkgs/python-modules/tvm-ffi/default.nix
@@ -0,0 +1,44 @@
+{
+  callPackage,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  cmake,
+  cython,
+  ninja,
+  python,
+  scikit-build-core,
+  setuptools-scm,
+}:
+
+buildPythonPackage rec {
+  pname = "tvm-ffi";
+  version = "0.1.9";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "apache";
+    repo = "tvm-ffi";
+    rev = "v${version}";
+    hash = "sha256-XnlM//WW2TbjbmzYBq6itJQ7R3J646UMVQUVhV5Afwc=";
+    fetchSubmodules = true;
+  };
+
+  build-system = [
+    cmake
+    cython
+    ninja
+    scikit-build-core
+    setuptools-scm
+  ];
+
+  dontUseCmakeConfigure = true;
+
+  postInstall = ''
+    ln -s $out/${python.sitePackages}/tvm_ffi/share $out/share
+  '';
+
+  passthru = callPackage ./variant.nix {
+    tvmFfiVersion = version;
+  };
+}
diff --git a/nix/pkgs/python-modules/tvm-ffi/variant.nix b/nix/pkgs/python-modules/tvm-ffi/variant.nix
@@ -0,0 +1,46 @@
+{
+  config,
+  cudaSupport ? config.cudaSupport,
+  metalSupport ? config.metalSupport or false,
+  rocmSupport ? config.rocmSupport,
+  xpuSupport ? config.xpuSupport or false,
+
+  cudaPackages,
+  rocmPackages,
+  xpuPackages,
+
+  lib,
+  stdenv,
+
+  tvmFfiVersion,
+}:
+
+let
+  flattenVersion =
+    version: lib.replaceStrings [ "." ] [ "" ] (lib.versions.majorMinor (lib.versions.pad 2 version));
+  backend =
+    if cudaSupport then
+      "cuda"
+    else if metalSupport then
+      "metal"
+    else if rocmSupport then
+      "rocm"
+    else if xpuSupport then
+      "xpu"
+    else
+      "cpu";
+  computeString =
+    if cudaSupport then
+      "cu${flattenVersion cudaPackages.cudaMajorMinorVersion}"
+    else if metalSupport then
+      "metal"
+    else if rocmSupport then
+      "rocm${flattenVersion (lib.versions.majorMinor rocmPackages.rocm.version)}"
+    else if xpuSupport then
+      "xpu${flattenVersion (lib.versions.majorMinor xpuPackages.oneapi-torch-dev.version)}"
+    else
+      "cpu";
+in
+{
+  variant = "tvm-ffi${flattenVersion (lib.versions.majorMinor tvmFfiVersion)}-${computeString}-${stdenv.hostPlatform.system}";
+}
