feat: extend kernel port tool operations by drbh · Pull Request #946 · huggingface/kernels

drbh · 2026-10-27T20:41:54Z

~~this pr adds a few kernel-port ops and post-port checks so recipes can handle native extensions and catch upstream drift at port time instead of at build or import time~~

this pr adds a few kernel-port ops and two new checks in the verify stage (which runs after the last op), so recipes can handle native extensions and catch upstream drift at port time instead of at build or import time

fix #includes when a glob move renames files (used in flash-attn3)
add rewrite, a structural find/replace (used in flash-attn3)
add extension to convert native extensions to ._ops (used in flash-attn3)
handle aliased and attribute-chain imports (not used in flash-attn3)
add mark_tests only= (used in flash-attn3)
add follow_includes=#true for kernel and manifest (used in flash-attn3 and manifest)
verify relative imports and torch-ext declarations (runs on every port, e.g. checks flash-attn3 and torch_src)
add expect blob= to pin overlay copies (not used in flash-attn3)
only match import edits at statement start (affects flash-attn3)

github-actions · 2026-10-27T21:20:44Z

Coverage report — `kernels/`

Measured on: Python 3.10 / Torch 2.13.0.
Other CI configurations are not included in this number.
Hardware-gated code paths (ROCm/XPU/NPU/Darwin/Windows) are excluded or unreachable on the Linux+CUDA runner.

Total coverage: 88.5% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	14	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	130	14	89%	53, 59-60, 63-64, 102, 165-170, 199, 219
src/kernels/archs.py	56	1	98%	94
src/kernels/backends.py	215	63	71%	42, 46, 50-53, 70, 92, 110, 119, 123, 127-129, 150, 159, 163, 167-169, 190, 201, 203, 210-213, 222, 226, 230-255, 263, 286-306
src/kernels/compat.py	9	1	89%	5
src/kernels/deps.py	70	1	99%	56
src/kernels/digest.py	58	5	91%	71-73, 136-137
src/kernels/hf_hub.py	62	2	97%	21, 23
src/kernels/importer.py	74	5	93%	158, 162, 165, 179-180
src/kernels/install.py	21	7	67%	76-100
src/kernels/layer/__init__.py	7	0	100%
src/kernels/layer/_interval_tree.py	103	4	96%	23, 52, 147, 150
src/kernels/layer/device.py	48	12	75%	42, 47-49, 91, 96-98, 101, 149, 152, 156
src/kernels/layer/func.py	93	5	95%	94, 119, 195, 315, 390
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/inspect.py	27	0	100%
src/kernels/layer/kernelize.py	83	7	92%	306, 314-315, 321, 325, 341-343
src/kernels/layer/layer.py	320	14	96%	212, 259, 447, 606-607, 628, 636, 647, 688, 692, 705, 760, 801, 876
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	136	36	74%	32, 35-42, 73, 86, 90, 99-100, 103, 106-109, 116-117, 120, 123-126, 133-134, 137, 140-143, 150-151, 154, 157-160, 235
src/kernels/load.py	71	2	97%	338, 378
src/kernels/locking.py	89	64	28%	35-83, 91-98, 102-125, 137, 152-159, 165-175, 179-186
src/kernels/python_deps.py	58	6	90%	59-60, 64-65, 101, 104
src/kernels/resolver.py	162	1	99%	248
src/kernels/status.py	50	2	96%	25, 79
src/kernels/validate.py	104	7	93%	9, 101, 172, 195-196, 232-233
src/kernels/variants.py	282	17	94%	65, 96, 117, 147, 256-257, 299-302, 304, 388-394, 400-406, 455-461
src/kernels/verify.py	102	6	94%	44, 176-178, 266-267
TOTAL	2469	283	89%

Updated by the Test kernels workflow on commit 201e204c22f0d1287531072f791c3bbb31e66224.

sayakpaul · 2026-10-28T05:30:59Z

I know this is still draft but would be nice to substantiate this with a real example.

this generates `flash-attn3/src` from `flash-attn3/port/port.kdl` and `flash-attn3/port/overlay`, pinned to the same upstream commit as the hand port (`0816ef12`) - the fixed `flash_attn_3` op namespace in `flash_api_stable.cpp` is rewritten by `extension` and the module entry point now lives in `torch-ext/torch_binding.cpp` (upstream's unused `PyInit__C` is kept) - `flash_api.cpp` is no longer vendored since it was never compiled - `flash_attn_interface.py` and the tests keep upstream's formatting and dispatch through `flash_attn_3_gpu = ops` - `build.toml` header lists come from `follow_includes` note: check ports runs kernel-port from the `kernel-port-declarative` branch until huggingface/kernels#946 is merged

drbh · 2026-10-28T16:53:28Z

just added an example that ports flash-attn3 huggingface/kernels-community#1283 it makes use of extension, rewrite, follow_includes=#true, mark_tests only= and the glob move that keeps #includes pointing at moved files, plus the post-port checks for relative imports and torch-ext declarations

this generates `flash-attn3/src` from `flash-attn3/port/port.kdl` and `flash-attn3/port/overlay`, pinned to the same upstream commit as the hand port (`0816ef12`) - the fixed `flash_attn_3` op namespace in `flash_api_stable.cpp` is rewritten by `extension` and the module entry point now lives in `torch-ext/torch_binding.cpp` (upstream's unused `PyInit__C` is kept) - `flash_api.cpp` is no longer vendored since it was never compiled - `flash_attn_interface.py` and the tests keep upstream's formatting and dispatch through `flash_attn_3_gpu = ops` - `build.toml` header lists come from `follow_includes` note: check ports runs kernel-port from the `kernel-port-declarative` branch until huggingface/kernels#946 is merged

Adds the `rewrite` op. A rewrite deletes whole lines, and splits identifiers before matching literal text, so a replacement cannot match inside a longer identifier.

…hain imports - Adds the `extension` op, deriving native extension conversions from a declaration. A declared module that upstream builds at import time with `var = torch.utils.cpp_extension.load(name="<module>", ...)` becomes `from ._ops import ops as var`, and the module-level statements that only fed the build (source paths, flags, the load import) are removed. Removing a trailing PYBIND11_MODULE no longer leaves blank lines at the end of the file. - remap_module also remaps `import a.b as x`. An unaliased `import a.b` binds `a`, so remapping it is an error rather than silently skipped. - convert_import converts an unaliased `import a.b.c` used as `a.b.c.attr` into `from a.b import c` with `c.attr` uses, refusing when `a` is used any other way or `c` is already a name in the module. - mark_tests only="<test>,..." marks just the named tests.

`kernel ... follow_includes=#true` (and `manifest ... torch_follow_includes=#true`) add every file the selected sources reach through #include, searched the way the compiler does: quoted includes next to the includer and then in the include directories, angle-bracket includes only in the include directories. A recipe then names compilation units instead of maintaining header lists that drift from what upstream actually includes.

- Relative imports must resolve: after the last op, every relative `from` import in the tree must name a module or package that is in the port (`torch-ext/<pkg>/_ops` is generated by the build). An upstream file that starts importing a module the recipe prunes now fails the port instead of failing at import time. - Every torch-ext declaration must match its definition: each function declared in a C-family file under torch-ext must have a definition in the output with the same parameter types (names and default arguments may differ). The error shows both lists and the first parameter that differs. - `expect blob=` pins the upstream file an overlay copy was made from. A find= guard only notices upstream changes to the lines it names, so an overlay copy of an upstream file could silently go stale. blob= pins the whole file by its git blob id; on a mismatch the error prints the git diff between the old and new blob, which is what to carry into the overlay copy. - Python import edits are spliced into the text so untouched code keeps its layout. Occurrences were found by text search alone, so an import that also appeared in a string literal or comment made the op refuse, and `import b` could match inside `from a import b`. Only occurrences that begin a statement, per libcst's tokenizer, now count, and every rewrite is located in the original text before any is applied.

…arch=#false A parenthesized from-import is rewritten like a one-line one unless a comment sits inside its parentheses, which the rewrite would drop. check_arch=#false passes check_arch=False to get_kernel, for kernels built only for devices the test runner may not have.

sayakpaul

So, from the PR description:

post-port checks so recipes can handle native extensions and catch upstream drift at port time instead of at build or import time

I think you meant the validation checks after the ops? Example:

if !mismatched.is_empty() {
        bail!(
            "verify: functions declared under torch-ext no longer match their \
             definitions (upstream changed a signature; update the binding):\n  {}",
            mismatched.join("\n  ")
        );
    }

drbh · 2026-11-10T15:57:30Z

yep thats what I meant, the verify stage that runs after the last op in run_pipeline_with. reworded the description to say that.

this pr adds two checks there: relative imports have to resolve to a module in the port (_ops excepted since the build generates it), and functions declared in c/c++/cuda files under torch-ext need a definition with the same param types. so upstream importing a module the recipe prunes, or changing a signature the binding depends on, fails the port instead of at import or link time

thanks for flagging!

drbh changed the title ~~Kernel port declarative~~ feat: extend kernel port tool operations Oct 28, 2026

drbh mentioned this pull request Oct 28, 2026

feat: generate flash-attn3 from a kernel-port recipe huggingface/kernels-community#1283

Open

drbh mentioned this pull request Nov 3, 2026

deep-ep: add DeepEP 2.5.0 (expert-parallel dispatch/combine) huggingface/kernels-community#1231

Open

3 of 6 tasks

drbh added 6 commits November 3, 2026 17:15

kernel-port: glob moves that keep #includes pointing at moved files

d784afb

kernel-port: rewrite, a structural find/replace for source code

3bd8ef8

Adds the `rewrite` op. A rewrite deletes whole lines, and splits identifiers before matching literal text, so a replacement cannot match inside a longer identifier.

kernel-port: extension redirects package-relative extension imports

f5dcf93

drbh marked this pull request as ready for review November 3, 2026 17:18

drbh force-pushed the kernel-port-declarative branch from b3c667e to f5dcf93 Compare November 3, 2026 17:18

sayakpaul approved these changes Nov 10, 2026

View reviewed changes

drbh merged commit adf4420 into main Nov 10, 2026
49 checks passed

Navigation Menu

Repository navigation

feat: extend kernel port tool operations - #946

feat: extend kernel port tool operations#946
drbh merged 7 commits into
mainfrom
kernel-port-declarative

drbh commented Oct 27, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 27, 2026 •

edited

Loading

Uh oh!

sayakpaul commented Oct 28, 2026

Uh oh!

drbh commented Oct 28, 2026

Uh oh!

sayakpaul left a comment

Uh oh!

drbh commented Nov 10, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

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

Uh oh!

Uh oh!

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

Uh oh!

Coverage report — kernels/

Uh oh!

sayakpaul commented Oct 28, 2026

Uh oh!

drbh commented Oct 28, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

drbh commented Nov 10, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

drbh commented Oct 27, 2026 •

edited

Loading

github-actions Bot commented Oct 27, 2026 •

edited

Loading

Coverage report — `kernels/`