feat: support rust gpu

docs/source/builder/writing-kernels.md

@@ -318,7 +318,7 @@ The following options can be set for a kernel:
 
 - `language` (required): the language and compute backend of the kernel,
   as `<language>-<backend>`. The currently supported values are `cpp-cpu`,
-  `cpp-cuda`, `cpp-metal`, `cpp-rocm`, `cpp-xpu`, and `rust-cpu`.
+  `cpp-cuda`, `cpp-metal`, `cpp-rocm`, `cpp-xpu`, `rust-cpu`, and `rust-cuda`.
   **The `cpu` backend is currently experimental and might still change.**
 - `depends` (required): a list of dependencies. The supported dependencies
   are listed in [`cpp-deps.nix`](https://github.com/huggingface/kernels/blob/main/builder/lib/cpp-deps.nix).
@@ -377,6 +377,31 @@ in Git so Nix can include them in the build. Rerun the command whenever Git
 dependencies in `Cargo.lock` change. Registry dependencies use the checksums in
 `Cargo.lock` and do not require this file.
 
+#### rust-cuda (experimental, and subject to change)
+
+Rust CUDA kernels are built with
+[cuda-oxide](https://github.com/NVIDIA/cuda-rust) and take the same options as
+`rust-cpu`, plus:
+
+- `cuda-capabilities` (optional): same as for `cuda`. Like with `nvcc`, the
+  kernels are compiled once per capability, to SASS, so nothing is JIT-compiled
+  at load time. Each pass sets `cfg(cuda_arch = "XX")` for its capability and
+  every lower one, so `#[cfg(cuda_arch = "90")]` works like
+  `__CUDA_ARCH__ >= 900`.
+
+Rust CUDA kernels require CUDA 13.0 or later, so they are only built for
+CUDA 13.0+ build variants (a lower `general.cuda.minver` is raised to 13.0). The CUDA
+bindings are generated with bindgen during the build, so building outside of
+Nix requires a CUDA 13 toolkit (found via `CUDA_TOOLKIT_PATH` or `CUDA_HOME`)
+and libclang.
+
+An extension can have only one `rust-cuda` component. Use the cuda-rust
+revision pinned in
+[`cuda-oxide`](https://github.com/huggingface/kernels/blob/main/nix-builder/pkgs/cuda-oxide/default.nix)
+and see
+[`relu-rust-cuda`](https://github.com/huggingface/kernels/tree/main/examples/kernels/relu-rust-cuda)
+for an example.
+
 ## Torch bindings
 
 ### Defining bindings

examples/kernels/flake.nix

@@ -121,6 +121,12 @@
           drv =
             sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cudaVersion}-${sys}"};
         }
+        {
+          name = "relu-rust-cuda-kernel";
+          path = ./relu-rust-cuda;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${cuda13Version}-${sys}"};
+        }
         {
           name = "relu-kernel-cpu";
           path = ./relu;

examples/kernels/relu-rust-cuda/CARD.md

@@ -0,0 +1,65 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+# If the org / user isn't a trusted publisher, pass `trust_remote_code=True` to the
+# `get_kernel` call. You can find whether this kernel is from a trusted publisher
+# by going to the kernel's Hub page and finding the "Trusted publisher" status at
+# the top of the page.
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Upstream
+
+The original source code for this kernel comes from {{ upstream }}.
+{% endif %}
+{% if source %}
+
+## Source
+
+The kernel-builder formatted source for this kernel is available at {{ source }}.
+{% endif %}