From e3dd69399216fa994bbbd7fbbf0a1602b61b62d4 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 9 Jul 2026 11:19:37 +0000 Subject: [PATCH 1/2] feat: add Helion as a supported Python dependency Helion kernels are pure-Python (Triton-based, JIT-compiled at runtime), so they build as noarch kernels. The pinned nixpkgs already ships python3Packages.helion (1.0.0, pytorch/helion), and it resolves torch through the Python package fixpoint, so it pairs with the builder's own torch variants. Verified end-to-end with a vector-add Helion kernel: builds as a torch-cuda noarch variant, passes get-kernel-check, records python-depends in metadata.json, and runs correctly through kernels.get_local_kernel on a GB10 GPU. Co-Authored-By: Claude Fable 5 --- docs/source/builder/writing-kernels.md | 4 ++-- docs/source/kernel-requirements.md | 1 + kernels-data/src/python_dependencies.json | 4 ++++ kernels/src/kernels/python_depends.json | 4 ++++ 4 files changed, 11 insertions(+), 2 deletions(-) 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" }] From 2839c9b08a5afc85629de024ce9e7f133b9c1181 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 9 Jul 2026 11:23:57 +0000 Subject: [PATCH 2/2] fix: unbreak Helion for non-CUDA build sets nixpkgs marks Helion as broken without cudaSupport, but Helion also supports ROCm and XPU through Triton, so noarch kernels declaring those backends failed to evaluate. Co-Authored-By: Claude Fable 5 --- nix-builder/overlay.nix | 10 ++++++++++ 1 file changed, 10 insertions(+) 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;