experiment: support rust kernels by drbh · Pull Request #724 · huggingface/kernels

drbh · 2026-07-20T18:00:08Z

this pr explores extending the kernel builder to support kernels written in rust. at the moment kernels are expected to be written a cpp or pure python. this pr updates the build process to:

add two new backends rust-cpu and rust-cuda
enable defining kernels written in rust in the build.toml
update cmake build step to compile rust with cargo
add cuda-oxide to rust build env

the goal of this experiment is to explore the recently released https://github.com/NVlabs/cuda-oxide library as a backend for hf kernels so developers would have the option to author kernels purely in rust.

additionally this change set is used in a small proof of concept rust impl of flash attn2 using cuda-oxide and tvmffi bindings here https://github.com/drbh/flash-attn2-oxide

the published result can be seen https://huggingface.co/kernels/drbh/flash-attn2-oxide and can be run in comparison to the reference kernel

# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "apache-tvm-ffi>=0.1.12",
#     "kernels>=0.16.0",
#     "numpy>=2.5.1",
#     "torch>=2.13.0",
# ]
# ///
import math

import torch
from kernels import get_kernel

oxide = get_kernel("drbh/flash-attn2-oxide", version=1, trust_remote_code=True)
reference = get_kernel("kernels-community/flash-attn2", version=2)

torch.manual_seed(0)
b, s, h, d = 2, 128, 8, 64
q = torch.randn(b, s, h, d, device="cuda", dtype=torch.float16)
k = torch.randn(b, s, h, d, device="cuda", dtype=torch.float16)
v = torch.randn(b, s, h, d, device="cuda", dtype=torch.float16)
scale = 1.0 / math.sqrt(d)

for causal in (False, True):
    # fwd is positional: (q, k, v, alibi_slopes, softmax_scale, is_causal,
    # window_size_left, window_size_right, softcap)
    out, lse = oxide.fwd(q, k, v, None, scale, causal, -1, -1, 0.0)
    ref = reference.flash_attn_func(
        q, k, v, dropout_p=0.0, softmax_scale=scale, causal=causal
    )
    ref = ref[0] if isinstance(ref, (tuple, list)) else ref

    max_diff = (out.float() - ref.float()).abs().max().item()
    ok = torch.allclose(out.float(), ref.float(), atol=2e-3, rtol=2.5e-2)
    print(f"causal={causal!s:<5} max|diff|={max_diff:.2e} {'ok' if ok else 'MISMATCH'}")
    assert ok

then run

uv run example/compare.py

and it outputs

causal=False max|diff|=4.88e-04 ok
causal=True  max|diff|=4.88e-04 ok

Warning

this work is an experiment, and may change or be closed at any point in time

HuggingFaceDocBuilderDev · 2026-07-20T18:02:01Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

github-actions · 2026-07-20T18:45:12Z

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: 84.3% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	13	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	63	7	89%	46, 49, 52-53, 56-57, 100
src/kernels/backends.py	212	62	71%	40, 44, 48-51, 68, 90, 108, 117, 121, 125-127, 148, 157, 161, 165-167, 188, 199, 201, 208-211, 224, 228, 232-252, 260, 283-303
src/kernels/compat.py	8	1	88%	5
src/kernels/deps.py	58	4	93%	59-60, 101, 104
src/kernels/hf_hub.py	62	5	92%	18, 20, 114, 136-137
src/kernels/importer.py	48	3	94%	106, 110, 113
src/kernels/install.py	50	5	90%	78, 116-117, 155, 173
src/kernels/layer/__init__.py	6	0	100%
src/kernels/layer/_interval_tree.py	103	4	96%	23, 52, 147, 150
src/kernels/layer/device.py	48	14	71%	42, 47-49, 91, 96-98, 101, 149, 152, 155-157
src/kernels/layer/func.py	82	7	91%	83, 113, 185, 303, 309, 322, 340
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/kernelize.py	74	8	89%	255, 281, 289-290, 296, 300, 316-318
src/kernels/layer/layer.py	211	16	92%	169, 212, 218, 231, 339, 419-420, 432, 441, 449, 460, 489, 493, 506, 559, 589
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	144	42	71%	27, 33, 36-43, 63-64, 70, 73-76, 90, 94, 103-104, 110, 113-116, 123-124, 130, 133-136, 143-144, 150, 153-156, 163-164, 170, 173-176, 257
src/kernels/load.py	74	9	88%	218, 224, 230-231, 252-264
src/kernels/locking.py	108	70	35%	42-109, 113-136, 144, 148-155, 159-169, 173-180
src/kernels/status.py	49	2	96%	23, 81
src/kernels/variants.py	278	22	92%	64, 95, 116, 146, 255-256, 298-301, 303, 387-394, 400-406, 437-443, 455-461, 611-613
src/kernels/verify.py	88	1	99%	32
TOTAL	1804	283	84%

Updated by the Test kernels workflow on commit d4bfd09481df7803c8dc31464720e81b4615b415.

danieldk · 2026-08-21T13:51:54Z

+The crate is built by cargo as a staticlib and whole-archive linked into the
+extension, so it must export the `__tvm_ffi_*` symbols itself. Rust kernels
+currently require the `[tvm-ffi]` framework, and the project root must include
+a `Cargo.lock`.


Just to be sure: Rust doesn't mark the linkage of any symbols to evade dlopen's RTLD_LOCAL right? I wouldn't expect so, but just to be sure.

danieldk · 2026-08-21T13:53:39Z

+depends = []
+src = ["Cargo.toml", "Cargo.lock", "src"]
+device-manifest = "kernels/Cargo.toml"
+ptx-dir = "kernels-ptx"


Is it worth exposing this? Maybe we should just use a standard directory?

danieldk · 2026-08-21T13:54:45Z

+dsl = "cuda-oxide"
+depends = []
+src = ["Cargo.toml", "Cargo.lock", "src"]
+device-manifest = "kernels/Cargo.toml"


What does device-manifest mean here?

danieldk · 2026-08-21T13:55:08Z

+backend = "cuda"
+dsl = "cuda-oxide"
+depends = []
+src = ["Cargo.toml", "Cargo.lock", "src"]


src? Don't we always explicitly specify all files? I'm surprised that this works?

danieldk · 2026-08-21T13:55:55Z

+    }
+}
+
+fn main() {}


I'm surprised that this is not built as a lib?

danieldk · 2026-08-21T14:29:03Z

+  cudaOxideSrc = pkgs.fetchFromGitHub {
+    owner = "drbh";
+    repo = "cuda-oxide";
+    rev = "bc5f33bb1d556671c35e475439bb821d25caa2cb";
+    hash = "sha256-ZIpS0XFDI9yDSXk33OyQ3mGB5gRhE7tldmmex226tLA=";
+  };


Should probably be in the overlay, maybe in an attr-set in there.

danieldk · 2026-08-21T14:30:39Z

+  cudaOxideRust =
+    if rust-bin != null then
+      rust-bin.fromRustupToolchainFile "${cudaOxideSrc}/rust-toolchain.toml"
+    else
+      rustc;
+  cudaOxideRustPlatform = pkgs.makeRustPlatform {
+    cargo = cudaOxideRust;
+    rustc = cudaOxideRust;
+  };
+  cudaOxideBackend = cudaOxideRustPlatform.buildRustPackage {
+    pname = "rustc-codegen-cuda";
+    version = "0.2.1";
+    src = cudaOxideSrc;
+    sourceRoot = "source/crates/rustc-codegen-cuda";
+    cargoLock = {
+      lockFile = "${cudaOxideSrc}/crates/rustc-codegen-cuda/Cargo.lock";
+      allowBuiltinFetchGit = true;
+    };
+    buildInputs = with pkgs; [
+      libffi
+      libxml2
+      zstd
+      zlib
+    ];
+    doCheck = false;
+  };


These should probably be factored out somewhere as an mk.* function.

Reading the rest of the file, maybe we should create a separate arch-rust.nix, at least for now?

danieldk · 2026-08-21T14:31:49Z

+  cargoDeps =
+    if hasRustKernels then
+      rustPlatform.importCargoLock {
+        lockFile = src + "/Cargo.lock";
+        allowBuiltinFetchGit = true;
+      }
+    else
+      null;


I don't think it's good to have this in the environment for C++ builds (even if set to null).

danieldk · 2026-08-21T14:32:17Z

+  preInstallCheck = lib.optionalString (hasRustKernels && cudaSupport) ''
+    addToSearchPath LD_LIBRARY_PATH "${lib.getOutput "stubs" cudaPackages.cuda_cudart}/lib/stubs"
+    export LD_LIBRARY_PATH
+  '';


LD_LIBRARY_PATH inside a derivation typically hides another issue.

danieldk · 2026-08-21T14:34:13Z

+      if rust-bin != null && builtins.pathExists (src + "/rust-toolchain.toml") then
+        [ (rust-bin.fromRustupToolchainFile (src + "/rust-toolchain.toml")) ]


Do we want to allow custom toolchains? Also, it seems that it can differ from cuda-oxide?

drbh temporarily deployed to testpypi August 3, 2026 13:49 — with GitHub Actions Inactive

drbh temporarily deployed to testpypi August 3, 2026 17:56 — with GitHub Actions Inactive

feat: support Rust kernels

76ad94d

drbh force-pushed the support-rust-kernels branch from 5e72710 to 76ad94d Compare August 19, 2026 01:21

drbh temporarily deployed to testpypi August 19, 2026 01:22 — with GitHub Actions Inactive

fix: load Rust kernels through kernels API

d4bfd09

drbh temporarily deployed to testpypi August 19, 2026 03:20 — with GitHub Actions Inactive

danieldk self-requested a review August 20, 2026 15:59

danieldk reviewed Aug 21, 2026

View reviewed changes

Navigation Menu

Uh oh!

experiment: support rust kernels - #724

experiment: support rust kernels#724
drbh wants to merge 2 commits into
mainfrom
support-rust-kernels

drbh commented Jul 20, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Jul 20, 2026

Uh oh!

github-actions Bot commented Jul 20, 2026 •

edited

Loading

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

danieldk Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		if rust-bin != null && builtins.pathExists (src + "/rust-toolchain.toml") then
		[ (rust-bin.fromRustupToolchainFile (src + "/rust-toolchain.toml")) ]

Uh oh!

Conversation

drbh commented Jul 20, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Jul 20, 2026

Uh oh!

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

Uh oh!

Coverage report — kernels/

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!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

github-actions Bot commented Jul 20, 2026 •

edited

Loading

Coverage report — `kernels/`