Add Helion as a supported Python dependency

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`
 

docs/source/kernel-requirements.md

@@ -200,6 +200,7 @@ The following dependencies are currently allowed:
 **General dependencies:**
 
 - `einops`
+- `helion`
 
 **Backend-specific dependencies:**
 

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" }]

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" }]

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;