experiment: support rust kernels

[drbh]

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]

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]

Coverage report — kernels/

Measured on: Python 3.10 / Torch 2.12.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.7% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	10	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	194	55	72%	40, 44, 48-51, 68, 90, 108, 117, 121, 125-127, 148, 170, 181, 188-191, 201, 205-225, 233, 256-276
src/kernels/compat.py	8	1	88%	5
src/kernels/deps.py	54	4	93%	58-59, 95, 98
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	81	7	91%	81, 111, 183, 301, 307, 320, 338
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	210	16	92%	167, 210, 216, 229, 337, 417-418, 430, 439, 447, 458, 487, 491, 504, 557, 587
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	130	34	74%	27, 33, 36-41, 61-62, 68, 71-74, 88, 92, 101-102, 108, 111-114, 121-122, 128, 131-134, 141-142, 148, 151-154, 235
src/kernels/lockfile.py	71	46	35%	37-104, 108-131
src/kernels/status.py	49	2	96%	23, 81
src/kernels/utils.py	313	55	82%	65, 77-81, 87-88, 246, 250, 253, 315, 323, 362-363, 401, 432, 437, 472, 701, 704, 706, 712, 725-726, 747-759, 763-770, 778, 782-792, 796-803, 841, 845, 864, 866
src/kernels/variants.py	262	19	93%	56, 87, 108, 138, 247-248, 289, 291, 371-378, 384-390, 421-427, 439-445, 534-536
src/kernels/verify.py	88	1	99%	32
TOTAL	1789	274	85%

Updated by the Test kernels workflow on commit 8237e30541d08878ada6046cc5200ec2b9f24b7e.