feat: prefer recipe ported einops by drbh · Pull Request #1092 · huggingface/kernels-community

drbh · 2026-08-21T14:43:44Z

*updated to match latest changeset

this pr is an example of using the kernel porting tool to port the upstream einops into the kernel builder format

first we clone the upstream source*

# first clone the repo locally
git clone https://github.com/arogozhnikov/einops.git /tmp/einops-src

# next we pin to a specific commit (same used in the recipe)
git -C /tmp/einops-src checkout ad2c8d67d08cbebb45e1aff2be21dc88635fc695

next we

# run the tool directly from the branch (could later be built into the cli tool)
nix run 'github:huggingface/kernels/pull/776/head#kernel-port' -- \
einops/port/port.kdl \ # the recipe to port
--dir /tmp/einops-src \ # the location of the source we cloned
--out ~/Projects/kernels-community/einops # the output location

running the command will show how each line of the recipe impacted the transformation and outputs all of the files added and deleted in the process.

[line   4] source              verified https://github.com/arogozhnikov/einops.git @ ad2c8d67d08c
[line   9] prune               removed 57 file(s), kept 27
[line  12] delete              removed 1 file(s)
[line  15] move                moved 8 file(s) to "tests"
[line  20] move                moved 18 file(s) to "torch-ext/einops"
[line  25] relativize_imports  rewrote 7 import(s) in 4 file(s)
[line  30] ensure_import       added 1 import(s) in 1 file(s)
[line  36] remap_module        rewrote 7 import(s) in 6 file(s)
[line  42] relativize_imports  rewrote 7 import(s) in 6 file(s)
[line  47] kernelize_imports   kernelized 26 import(s) in 8 file(s)
[line  54] replace             1 replacement(s) in 1 file(s)
[line  59] replace             1 replacement(s) in 1 file(s)
[line  64] replace             7 replacement(s) in 7 file(s)
[line  70] replace             1 replacement(s) in 1 file(s)
[line  75] manifest            wrote build.toml (noarch)

check-ports.yaml reruns the port in ci and fails if einops/src does not match the output, so the generated tree is still reviewable but cannot drift from the recipe. prs only check the ports they touch

note: flake.nix is not generated so it stays at einops/flake.nix and points nix at ./src, build.yaml already handles this layout (#1109)

drbh · 2026-08-21T14:49:24Z

for completeness, we can bump the commit hash from ad2c8d67d08cbebb45e1aff2be21dc88635fc695 to the latest head e0d5eb4fd535945ff65d309206d0f1754a926821

and trying to rerun the command shows

nix run 'github:huggingface/kernels/pull/776/head#kernel-port' -- .ports/einops/port.kdl --dir /tmp/einops-src --out einops
[line  14] source              verified https://github.com/arogozhnikov/einops.git @ e0d5eb4fd535
[line  18] prune               removed 56 file(s), kept 27
[line  19] delete              removed 10 file(s)
[line  21] move                moved 17 file(s) to "torch-ext/einops"
error: recipe line 25: relativize_imports: expected exactly 7 change(s) but made 13 - upstream drifted; review the new rewrites and update changes=

at this point I simply asked and llm to fix the recipe so the port works correctly and copy pasted the command and error output above.

the exact output of this process can be seen in #1093

sayakpaul · 2026-08-25T08:28:17Z

I think we should have a single standard to start this. I would prefer if we required users to specify the source and the commit in the recipe itself instead of having them locally clone and checkout to a particular branch first. Or, that should be a better default behaviour, IMO.

9 of the 15 files in the overlay are tests we added that are not ported from upstream

We should port tests, too. Then it will be down on us to make modifications there e.g., for running specific tests on our CI, etc. Cc: @danieldk here as well.

sayakpaul · 2026-08-25T08:30:15Z

Going to run the commands from the PR description and share my findings here as well.

sayakpaul · 2026-08-25T08:33:23Z

+    find="from .packing import pack, unpack\n" \
+    with="from .packing import pack, unpack\nfrom . import array_api\n"
+
+// The rewritten tests and the flake.


We should not have to rewrite the tests much though. I think in the tests we would like to ensure that they use get_kernel() which already seems to be the case. So, I think we should rather provide similar instructions in the recipe to account for this.

pytest.mark.kernels_ci should probably be left to us because it's better if we decided which tests are lightweight and important enough to run on our CI.

And also, I think tests should not be overlays no? Currently, we include them inside of the actual kernel source.

the tests actually come from upstream and the port was updated so we rewrite the import to use get_kernels, this added some more lines to the port - but now we do not have any tests in the overlay file.

however if we were to add tests that were not originally upstream and rewritable - they would belong in the overlays folder - since that should contains all of the code that is not upstream. the "kernel source" which is now in the einops/src file should not contain any code that is not output from the port too. we keep this clear separation so the einops/src only contains generated code (we never mix non generated code with the tools output)

I agree. Those tests (that we modify heavily) should probably be reserved for just CI-level testing. Those are fine to have in overlays.

sayakpaul · 2026-08-25T09:43:45Z

#1100 -- used the tool on flash-attn-ops. Major changes are in a layer norm file and I provided some comments there.

drbh · 2026-08-27T00:02:19Z

/kernel-bot build-and-stage einops

github-actions · 2026-08-27T00:02:36Z

Build request processed.

Command: /kernel-bot build-and-stage einops
Mode: build and stage
Target branch: pr-1092
PR head SHA: 26b94d072a85a3609dea42a667ffe368126ddfc8
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (2):

einops (build-mac.yaml): https://github.com/huggingface/kernels-community/actions/runs/33025463511
einops (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/33025465620

sayakpaul · 2026-08-27T02:05:50Z

@drbh can we build using the bot without the PR not having landed in main?

danieldk · 2026-08-27T13:13:48Z

+
+replace in="tests/test_einsum.py" count=1 \
+    find=#"""
+from einops.tests import collect_test_backends


This seems quite verbose? It would be nicer to have something less literal and more declarative that is "for every import of the kernel of the shape import kernel replace by kernel = get_kernel(...) and every import of the shape from kernel import foo replace by foo = kernel.foo. All these full snippets are going to be very fragile for active upstream projects.

drbh · 2026-08-27T19:36:11Z

/kernel-bot build-and-stage einops

github-actions · 2026-08-27T19:36:34Z

Build request processed.

Command: /kernel-bot build-and-stage einops
Mode: build and stage
Target branch: pr-1092
PR head SHA: 2f5d4fe7d964a2ed5a7d55b89041dc2be17966e4
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (2):

einops (build-mac.yaml): https://github.com/huggingface/kernels-community/actions/runs/33109260091
einops (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/33109261646

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-staging/einops

sayakpaul

Thanks, left some comments. I like the idea of validating a port, really cool!

sayakpaul · 2026-09-01T10:08:38Z

+# A ported kernel's <kernel>/src is generated from <kernel>/port/port.kdl and
+# committed. Regenerate it here and fail if it drifted, so that a stale overlay
+# silently reverting an in-tree fix, a hand-edited build.toml field the manifest
+# op drops, or an overlay file edited without its generated twin are all loud.


What do we consider to be a drift here?

drift here is just einops/src not matching what the recipe outputs. the cases this catches are an overlay file going stale and undoing a fix that was made directly in src, a manual edit to build.toml that the manifest op drops on the next run, and the recipe being edited without rerunning the port (port-provenance.json stores the recipe hash so that one shows up as a diff in the provenance file). updated the comment in the workflow to list these

sayakpaul · 2026-09-01T10:09:21Z

+    types: [opened, synchronize, reopened]
+    paths:
+      - "*/port/**"
+      - "*/src/**"


So for kernels containing valid porting recipes, we expect their sources to be present under src, which is as expected 👍

sayakpaul · 2026-09-01T10:10:07Z

+        env:
+          # kernel-port still lives on the kernel-port-tool branch; point this at
+          # github:huggingface/kernels#kernel-port once that branch merges.
+          KERNEL_PORT: "nix run github:huggingface/kernels/kernel-port-tool#kernel-port --"


Let's add a TODO here to update when the PR is eventually merged.

Also, once this is shipped do we expect users to run with nix run or through the kernel-builder CLI?

Let's add a TODO here to update when the PR is eventually merged.

done, also pinned the runner to a commit on the kernel-port-tool branch instead of the branch head so changes there dont affect open prs

Also, once this is shipped do we expect users to run with nix run or through the kernel-builder CLI?

nix run for now since the tool only exists on the branch. once it merges I think we can expose it as a kernel-builder subcommand and update KERNEL_PORT in the workflow, the script just reads that env var so either works

sayakpaul · 2026-09-01T10:13:12Z

@@ -0,0 +1,12 @@
+{
+  "format": 1,


Self TODO: review the provenance generation logic.

sayakpaul · 2026-09-01T10:14:53Z

@@ -0,0 +1,80 @@
+recipe version=1


Okay so IIUC we didn't need any overlay in this case? I thought the flake.nix would go to the overlay directory? How are we handling that?

yea no overlay is needed for einops since everything comes from upstream. flake.nix is not generated so it stays at einops/flake.nix with path = ./src and only the tool output lives under src/. build.yaml already picks up <kernel>/src/build.toml when it exists

flake.nix is not generated so it stays at einops/flake.nix

Oh okay. Earlier, I think we wanted to keep the external files in the overlay. What is the general guidelines behind putting stuff under an overlay folder then?

this pr fixes the macos ci failures where `get-kernel-check` rejects a kernel that built fine. first hit on the einops port in huggingface/kernels-community#1092 tldr; the check handed `$out` to the resolver and let it pick a variant by detecting the backend at runtime. metal is the only backend detected from live hardware (`torch.backends.mps.is_available()`, everything else is baked into the torch build). the sandboxed macos build cannot see the gpu, so the `torch-metal` variant is never selected and the check fails even though the artifact is fine repro (no kernel build needed) ```bash mkdir -p /tmp/fake-kernel/torch-metal && touch /tmp/fake-kernel/torch-metal/metadata.json kernels/.venv/bin/python -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('/tmp/fake-kernel'))" ``` ``` ValueError: Cannot parse metadata from `"/tmp/fake-kernel/torch-metal/metadata.json"`: EOF while parsing a value at line 1 column 0 ``` on a mac with a visible gpu the variant is selected fine (it gets as far as reading the empty `metadata.json`). same command inside a seatbelt sandbox ```bash sandbox-exec -p '(version 1)(allow default)(deny iokit-open)' \ kernels/.venv/bin/python -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('/tmp/fake-kernel'))" ``` ``` FileNotFoundError: Cannot find a build variant for this system in /tmp/fake-kernel: ``` `sandbox-exec` is the same seatbelt mechanism the nix darwin sandbox uses. denying iokit hides the gpu, `mps.is_available()` returns false, and the same directory no longer resolves note* this is also why only some metal builds fail in ci. arch metal builds set `__noChroot = metalSupport` since they need the host metal toolchain, so the gpu is still visible during the check. noarch metal builds like einops run fully sandboxed and always fail. the first command also fails on linux since there is no mps at all the check now globs `$out` for `*/metadata.json` and points the resolver at that exact variant dir instead of auto detecting from the build host. each extension derivation produces exactly one variant, so zero or multiple variants still fail the check

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

…e concise

github-actions Bot added the feature New functionality / capability label Aug 21, 2026

drbh mentioned this pull request Aug 21, 2026

feat: experimental kernel porting tool huggingface/kernels#776

Open

sayakpaul reviewed Aug 25, 2026

View reviewed changes

Comment thread einops/src/port-provenance.json

sayakpaul reviewed Aug 25, 2026

View reviewed changes

Comment thread .ports/einops/port.kdl Outdated

sayakpaul reviewed Aug 25, 2026

View reviewed changes

Comment thread .ports/einops/overlay/tests/test_other.py

sayakpaul reviewed Aug 25, 2026

View reviewed changes

Comment thread .ports/einops/port.kdl Outdated

sayakpaul mentioned this pull request Aug 25, 2026

[don't merge] port flash-attn-ops using the porting tool #1100

Closed

drbh force-pushed the port-einops branch from 8fc092a to 926f8e0 Compare August 26, 2026 23:30

drbh force-pushed the port-einops branch from 26b94d0 to cf5474a Compare August 27, 2026 00:18

drbh changed the base branch from main to nested-build-toml-support August 27, 2026 00:20

danieldk reviewed Aug 27, 2026

View reviewed changes

Base automatically changed from nested-build-toml-support to main August 27, 2026 14:24

drbh force-pushed the port-einops branch from cf5474a to 29754fe Compare August 27, 2026 14:24

drbh mentioned this pull request Aug 31, 2026

fix: point get-kernel-check at the built variant huggingface/kernels#802

Merged

sayakpaul reviewed Sep 1, 2026

View reviewed changes

drbh added 4 commits September 2, 2026 15:30

feat: prefer recipe ported einops

8164277

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

feat: move einops kernel sources under src/

807a2da

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

feat: generate einops src/ from a port recipe

42012e2

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

feat: use multi-line strings in the einops port recipe

17a38d6

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

drbh added 5 commits September 2, 2026 15:30

feat: move the einops flake out of the port overlay

989f395

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

feat: verify ported kernels regenerate from their recipe in CI

f29c62e

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

fix: point the port check at the kernel-port-tool branch

0b0ef61

Signed-off-by: David Holtz <david.richard.holtz@gmail.com>

feat: use ensure_import and kernelize_imports to make recipe much mor…

333e069

…e concise

fix: improve port ci checks and fix typos

0b35517

drbh force-pushed the port-einops branch from ce69618 to 0b35517 Compare September 2, 2026 19:30

Uh oh!

Conversation

drbh commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

drbh commented Aug 21, 2026

Uh oh!

sayakpaul commented Aug 25, 2026

Uh oh!

sayakpaul commented Aug 25, 2026

Uh oh!

Uh oh!

sayakpaul Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sayakpaul commented Aug 25, 2026

Uh oh!

drbh commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sayakpaul commented Aug 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

drbh commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

drbh commented Aug 21, 2026 •

edited

Loading

sayakpaul Aug 25, 2026 •

edited

Loading

github-actions Bot commented Aug 27, 2026 •

edited

Loading

github-actions Bot commented Aug 27, 2026 •

edited

Loading