nix-builder: redesign kernel testing

docs/source/builder/agents-guide.md

@@ -1,6 +1,6 @@
 # Develop kernels with agents
 
-Code agents are a good fit to build custom kernels because the hard part is not just writing in Domain Specific Language (DSLs) like CUDA. You also need the right project layout, PyTorch bindings, architecture-specific choices, model-specific integration, and trustworthy benchmarks. 
+Code agents are a good fit to build custom kernels because the hard part is not just writing in Domain Specific Language (DSLs) like CUDA. You also need the right project layout, PyTorch bindings, architecture-specific choices, model-specific integration, and trustworthy benchmarks.
 
 Kernels on Hugging Face are compatible with agents via skills and the `hf` CLI. The `cuda-kernels`, `rocm-kernels`, `xpu-kernels`, and `cpu-kernels` skills contain knowledge so an agent can generate and publish a complete kernel project, instead of isolated snippets.
 
@@ -10,7 +10,7 @@ This guide is for **authoring new kernels**. If you only want to **load an exist
 
 You need:
 
-- a coding agent that supports skills, such as Claude Code, Codex, Cursor, or OpenCode  
+- a coding agent that supports skills, such as Claude Code, Codex, Cursor, or OpenCode
 - a clear target: library, model, operation, GPU, dtype, and representative shapes
 
 The skill currently focuses on NVIDIA GPUs such as **H100**, **A100**, and **T4**, and on integration patterns for **transformers** and **diffusers**.
@@ -31,14 +31,14 @@ kernel-builder skills add --claude
 
 Writing kernels is a hard problem, so be specific to agents. A robust prompt will declare all core attributes, including:
 
-- the library, for example `transformers` or `diffusers`  
-- the model id, for example `Qwen3-8B` or `LTX-Video`  
-- the operation, for example `RMSNorm`, attention, RoPE, `GEGLU`, or `AdaLN`  
-- the target GPU, for example `H100`, `A100`, or `T4`  
-- the dtype, for example `bfloat16`, `float16`, or `float32`  
+- the library, for example `transformers` or `diffusers`
+- the model id, for example `Qwen3-8B` or `LTX-Video`
+- the operation, for example `RMSNorm`, attention, RoPE, `GEGLU`, or `AdaLN`
+- the target GPU, for example `H100`, `A100`, or `T4`
+- the dtype, for example `bfloat16`, `float16`, or `float32`
 - the outputs you expect: kernel code, bindings, tests, and benchmarks
 
-In practice, you can often skip some of these and the agent will infer based on common practice, but if you know a detail declare it. 
+In practice, you can often skip some of these and the agent will infer based on common practice, but if you know a detail declare it.
 
 For example:
 
@@ -69,7 +69,7 @@ examples/your_model/
 │   └── torch_binding.cpp       # PyTorch C++ bindings
 ├── benchmark_rmsnorm.py        # Micro-benchmark script
 ├── build.toml                  # kernel-builder config
-├── setup.py                    # pip install -e .
+├── setup.py                    # python setup.py build_kernel
 └── pyproject.toml
 ```
 
@@ -105,15 +105,15 @@ cuda-capabilities = ["9.0"]  # H100
 
 First check that:
 
-- `backends = ["cuda"]` is correct for your project  
-- the kernel source files are listed correctly  
-- the Torch binding sources are included under `[torch]`  
+- `backends = ["cuda"]` is correct for your project
+- the kernel source files are listed correctly
+- the Torch binding sources are included under `[torch]`
 - `cuda-capabilities` is only set when the kernel truly targets specific architectures
 
 For architecture-specific kernels, typical capability values are:
 
-- H100: `9.0`  
-- A100: `8.0`  
+- H100: `9.0`
+- A100: `8.0`
 - T4: `7.5`
 
 If the kernel does **not** require a specific capability, the kernels docs recommend leaving `cuda-capabilities` unset so the builder can target all supported capabilities. In practice, you can prompt your agent to review the `build.toml` for excessive definitions. Agents have a tendency to over-specify capabilities.
@@ -126,17 +126,17 @@ The kernel should be registered as Torch ops in `torch-ext/torch_binding.cpp`, w
 
 Make sure the integration matches the library:
 
-- **transformers**: patch the target modules directly, often RMSNorm modules whose class name contains `RMSNorm`  
+- **transformers**: patch the target modules directly, often RMSNorm modules whose class name contains `RMSNorm`
 - **diffusers**: inspect the actual pipeline structure before patching, because modules and attention processors can differ across pipelines
 
 > [!NOTE]
 > One common issue is that the agent will not integrate the kernel at all. Typically because the project's context is so long.
 
 A few patterns matter in practice for the integration code:
 
-- In **transformers**, RMSNorm modules generally have weights, but epsilon may be exposed as `variance_epsilon` or `eps` depending on the model.  
-- In **diffusers**, some RMSNorm modules may have `weight=None`, so the integration code needs to handle both weighted and unweighted cases.  
-- In **diffusers**, checking `type(module).__name__` is often more reliable than `isinstance(...)` for matching RMSNorm modules across implementations.  
+- In **transformers**, RMSNorm modules generally have weights, but epsilon may be exposed as `variance_epsilon` or `eps` depending on the model.
+- In **diffusers**, some RMSNorm modules may have `weight=None`, so the integration code needs to handle both weighted and unweighted cases.
+- In **diffusers**, checking `type(module).__name__` is often more reliable than `isinstance(...)` for matching RMSNorm modules across implementations.
 - If a diffusers pipeline uses CPU offloading, inject custom kernels **before** enabling offload.
 
 For attention, prefer the model library's existing optimized path when one already exists. For example, in `transformers`, Flash Attention 2 is usually the right baseline for attention, while custom kernels are especially useful for operations like RMSNorm and other targeted hotspots.
@@ -160,22 +160,22 @@ nix run nixpkgs#cachix -- use huggingface
 
 There are two main benchmarks to consider:
 
-1. an isolated kernel micro-benchmark  
+1. an isolated kernel micro-benchmark
 2. an end-to-end benchmark in the real model or pipeline
 
 The agent will generate both benchmarks based on the agent skills examples. Typically as a script called `benchmark_example.py`. If you have access to the target hardware, you can run it to verify the kernel works. For example, the agent will generat a table like this:
 
 ```markdown
-| Shape | Custom (ms) | PyTorch (ms) | Speedup |
-| :---- | :---: | :---: | :---: |
-| [1x128x4096] | 0.040 | 0.062 | **1.58x** |
-| [1x512x4096] | 0.038 | 0.064 | **1.69x** |
-| [1x1024x4096] | 0.037 | 0.071 | **1.90x** |
-| [1x2048x4096] | 0.045 | 0.091 | **2.03x** |
-| [1x4096x4096] | 0.071 | 0.150 | **2.12x** |
-| [4x512x4096] | 0.056 | 0.093 | **1.67x** |
-| [8x256x4096] | 0.045 | 0.092 | **2.06x** |
-| [1x8192x4096] | 0.109 | 0.269 | **2.47x** |
+| Shape         | Custom (ms) | PyTorch (ms) |  Speedup  |
+| :------------ | :---------: | :----------: | :-------: |
+| [1x128x4096]  |    0.040    |    0.062     | **1.58x** |
+| [1x512x4096]  |    0.038    |    0.064     | **1.69x** |
+| [1x1024x4096] |    0.037    |    0.071     | **1.90x** |
+| [1x2048x4096] |    0.045    |    0.091     | **2.03x** |
+| [1x4096x4096] |    0.071    |    0.150     | **2.12x** |
+| [4x512x4096]  |    0.056    |    0.093     | **1.67x** |
+| [8x256x4096]  |    0.045    |    0.092     | **2.06x** |
+| [1x8192x4096] |    0.109    |    0.269     | **2.47x** |
 ```
 
 Interpret the results carefully. A kernel can show a large isolated speedup but only a modest end-to-end gain if that operation is a small fraction of total runtime. In the LTX-Video example from [the blog we wrote](https://huggingface.co/blog/custom-cuda-kernels-agent-skills), the generated RMSNorm kernel improved the isolated benchmark by about **1.88x** on average, but end-to-end video generation improved by about **6%**, which matched the fact that RMSNorm accounted for only a small share of total compute.
@@ -186,7 +186,7 @@ Once the project is correct and benchmarked, you can build Hub-compatible artifa
 
 ```shell
 # install the hf CLI tool
-hf skills add 
+hf skills add
 
 # Authenticate
 hf auth login
@@ -216,4 +216,5 @@ from kernels import get_kernel
 kernel = get_kernel("your-org/your-kernel", version=1)
 ```
 
-Well done! You have now built a custom kernel and published it to the Hub.
+Well done! You have now built a custom kernel and published it to the Hub.
+

docs/source/builder/build.md

@@ -69,25 +69,27 @@ for monitoring the build. The compiled kernel will then be in the local
 
 `kernel-builder` provides shells for developing kernels. In such a shell,
 all required dependencies are available, as well as `kernel-builder` for generating
-project files. For example:
+project files. For example, you can use the development shell to build a
+arch (AOT-compiled) kernel:
 
 ```bash
 $ kernel-builder devshell
 # A devshell is opened in which you can run the following commands:
 $ kernel-builder create-pyproject
 $ cmake -B build-ext
-$ cmake --build build-ext
+$ cmake --build build-ext --target local_install
 ```
 
-If you want to test the kernel as a Python package, you can do so.
-`kernel-builder devshell` will automatically create a virtual environment in
-the `.venv` and activate it. You can install the kernel as a regular
-Python package in this virtual environment:
+This will build the kernel and puts the output in the `build` directory
+and can be used with the `kernels` library.
+
+Noarch (JIT-compiled) kernels do not use CMake. For this reason, we also
+create a `setup.py` that works both for arch and noarch kernels:
 
 ```bash
 $ kernel-builder devshell
 $ kernel-builder create-pyproject
-$ pip install --no-build-isolation -e .
+$ python setup.py build_kernel
 ```
 
 Development shells are available for every build configuration. For

docs/source/builder/local-dev.md

@@ -27,12 +27,24 @@ $ kernel-builder create-pyproject -f
 The `-f` flag is optional and instructs `kernel-builder` to overwrite
 existing files.
 
-It is recommended to do an editable install of the generated project into
-your Python virtual environment for development:
+You can build the kernel with
 
 ```bash
-$ pip install wheel # Needed once to enable bdist_wheel.
-$ pip install --no-build-isolation -e .
+$ python setup.py build_kernel
+```
+
+This builds the kernel and puts the variant that is compatible with the build
+environment in `build`. The build can then be loaded directly with `kernels`:
+
+```shell
+$ python -c 'import pathlib; import kernels; k = kernels.get_local_kernel(pathlib.Path("build")); print(k)'
+```
+
+For AOT kernels, if you want to skip the CMake configuration step in subsequent
+builds, you can also run Ninja directly to do incremental builds:
+
+```shell
+$ ninja -C _cmake_build local_install
 ```
 
 You can also create a Python project for a kernel in another directory:

docs/source/builder/writing-kernels.md

@@ -409,10 +409,35 @@ def relu_fwd_fake(input: torch.Tensor) -> torch.Tensor:
 
 ## Kernel tests
 
-Kernel tests are stored in the `tests` directory. Since running all
-kernel tests in CI may be prohibitively expensive, the `pyproject.toml`
-generated by the builder adds support for the special `kernels_ci`
-PyTest marker that can be used as follows:
+### Use `get_kernel` in tests
+
+Kernel tests are stored in the `tests` directory. Tests must not use direct
+imports, but instead use `get_kernel` to test the kernel as it will be used.
+For example:
+
+```python
+import kernels
+import torch
+import torch.nn.functional as F
+
+relu = kernels.get_kernel("kernels-community/relu", version=1)
+
+def test_relu():
+    x = torch.randn(1024, 1024, dtype=torch.float32, device=torch.device("cuda"))
+    y = relu.relu(x, torch.empty_like(x))
+    y_ref = F.relu(x)
+    torch.testing.assert_close(y_ref, y)
+```
+
+Development shells (`kernel-builder devshell`/`kernel-builder testshell`)
+will set the `LOCAL_KERNELS` variable to ensure that the kernel will be
+loaded from the development environment.
+
+### Mark CI tests
+
+Since running all kernel tests in CI may be prohibitively expensive, the
+`pyproject.toml` generated by the builder adds support for the special
+`kernels_ci` PyTest marker that can be used as follows:
 
 ```python
 import pytest