diff --git a/docs/source/builder/writing-kernels.md b/docs/source/builder/writing-kernels.md index 3f3d7d7d..9b1b1a07 100644 --- a/docs/source/builder/writing-kernels.md +++ b/docs/source/builder/writing-kernels.md @@ -205,8 +205,8 @@ The following sections enumerate all supported options for `build.toml`. formatted source repository. This repository must contain a `build.toml` and `flake.nix` so that it can be pulled and built with the kernel builder. - `python-depends` (**experimental**): a list of additional Python dependencies - that the kernel requires. The only supported dependencies are `einops` - and `nvidia-cutlass-dsl`. + that the kernel requires. The only supported dependencies are `einops`, + `helion`, and `nvidia-cutlass-dsl`. ### `general.hub` diff --git a/docs/source/kernel-requirements.md b/docs/source/kernel-requirements.md index f05beacc..eb432481 100644 --- a/docs/source/kernel-requirements.md +++ b/docs/source/kernel-requirements.md @@ -200,6 +200,7 @@ The following dependencies are currently allowed: **General dependencies:** - `einops` +- `helion` **Backend-specific dependencies:** diff --git a/kernels-data/src/python_dependencies.json b/kernels-data/src/python_dependencies.json index e46c200c..266904e8 100644 --- a/kernels-data/src/python_dependencies.json +++ b/kernels-data/src/python_dependencies.json @@ -4,6 +4,10 @@ "nix": ["einops"], "python": [{ "pkg": "einops", "import": "einops" }] }, + "helion": { + "nix": ["helion"], + "python": [{ "pkg": "helion", "import": "helion" }] + }, "tvm-ffi": { "nix": ["tvm-ffi"], "python": [{ "pkg": "apache-tvm-ffi", "import": "tvm_ffi" }] diff --git a/kernels/src/kernels/python_depends.json b/kernels/src/kernels/python_depends.json index e46c200c..266904e8 100644 --- a/kernels/src/kernels/python_depends.json +++ b/kernels/src/kernels/python_depends.json @@ -4,6 +4,10 @@ "nix": ["einops"], "python": [{ "pkg": "einops", "import": "einops" }] }, + "helion": { + "nix": ["helion"], + "python": [{ "pkg": "helion", "import": "helion" }] + }, "tvm-ffi": { "nix": ["tvm-ffi"], "python": [{ "pkg": "apache-tvm-ffi", "import": "tvm_ffi" }] diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index f5bcb994..f9480651 100644 --- a/nix-builder/overlay.nix +++ b/nix-builder/overlay.nix @@ -147,6 +147,16 @@ final: prev: }; }); + helion = python-super.helion.overrideAttrs ( + _: prevAttrs: { + # nixpkgs marks Helion as broken without CUDA, but it also + # supports ROCm and XPU through Triton. + meta = prevAttrs.meta // { + broken = false; + }; + } + ); + jax = python-super.jax.overrideAttrs ( _: prevAttrs: { dontUsePytestCheck = true;