cuda-kernels

# CUDA Kernels for Diffusers & Transformers This skill provides patterns and guidance for developing optimized CUDA kernels targeting NVIDIA GPUs (H100, A100, T4) for use with HuggingFace **diffusers** and **transformers** libraries. ## Hard Constraints — Read Before Writing Any Code Kernels MUST build with [kernel-builder](https://github.com/huggingface/kernels) and meet the [Kernel Hub requirements](https://huggingface.co/docs/kernels/kernel-requirements). kernel-builder compiles against the **Python limited API (ABI3)** so a single binary works for Python 3.9+ across versions. Several patterns that are standard in generic PyTorch-extension tutorials are therefore **hard build failures** here. Do not use them, even if PyTorch documentation or your training data suggests them. ### Disallowed patterns — never generate these | ❌ Never use | Why it fails | ✅ Use instead | |---|---|---| | pybind11 in any form: `#include <torch/extension.h>`, `#include <pybind11/...>`, `PYBIND11_MODULE(...)`, `py::arg`, any `py::` symbol | pybind11 is incompatible with the limited API (ABI3); the build does not compile | `TORCH_LIBRARY_EXPAND` in `torch-ext/torch_binding.cpp` (see below). Note: `torch/extension.h` transitively includes pybind11 — include `torch/torch.h` + `torch/library.h` instead | | Hand-written `setup.py` / `pyproject.toml` using `torch.utils.cpp_extension` (`CUDAExtension`, `BuildExtension`, `cpp_extension.load`, `load_inline`) | setuptools extensions are not ABI3 and bypass `build.toml`; kernel-builder owns the build | `build.toml` + `nix run .#build-and-copy -L`. For an editable dev install, generate the project files with `kernel-builder create-pyproject -f` — never write them by hand | | `TORCH_LIBRARY(my_kernel, m)`, `TORCH_LIBRARY_FRAGMENT(...)`, or `TORCH_LIBRARY_IMPL(...)` with a hardcoded namespace | kernel-builder suffixes the op namespace with a per-build hash (e.g. `_my_kernel_a1b2c3d`); a hardcoded name never resolves | `TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops)` from the generated `registration.h` | | Hardcoded `torch.ops.my_kernel.fn(...)` calls in Python | Same namespace mangling — the op namespace name is only known at build time | `from ._ops import ops` then `ops.fn(...)` | | Hand-written `PyMODINIT_FUNC PyInit__...` or any manual CPython module init | Generated by `REGISTER_EXTENSION`; duplicating it breaks module loading | `REGISTER_EXTENSION(TORCH_EXTENSION_NAME)` exactly once, in `torch_binding.cpp` | | Non-limited CPython API calls (`PyArg_ParseTuple`, direct `PyObject*` manipulation) | Violates ABI3 | Stay within the torch C++ API: `torch::Tensor`, `TORCH_CHECK`, `at::cuda::*` | | Absolute imports of your own package inside `torch-ext/` (`from my_kernel.utils import x`) | The package directory is renamed when loaded from the Hub; absolute imports break | Relative imports only: `from .utils import x`, `from ._ops import ops` | | Runtime Python deps beyond `torch` (and `einops` if truly needed) | Hub compliance restricts kernel dependencies; imports of numpy, triton, packaging, etc. are rejected | Standard library + `torch` only | | Python-side `@torch.library.custom_op` as the primary binding | The op must be registered in C++ so it ships in the compiled extension | C++ registration via `TORCH_LIBRARY_EXPAND`; Python-side `torch.library.register_fake` is only for adding a fake/meta impl (see torch.compile section) | ### The only supported binding pattern `registration.h` and `_ops.py` are **generated by kernel-builder** — reference them, never write them yourself. **`torch-ext/torch_binding.h`:** ```cpp #pragma once #include <torch/torch.h> void my_kernel_forward(torch::Tensor &out, torch::Tensor const &input); ``` **`torch-ext/torch_binding.cpp`:** ```cpp #include <torch/library.h> #include "registration.h" #include "torch_binding.h" TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("my_kernel_forward(Tensor! out, Tensor input) -> ()"); ops.impl("my_kernel_forward", torch::kCUDA, &my_kernel_forward); } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) ``` **`torch-ext/my_kernel/__init__.py`:** ```python import torch from ._ops import ops def my_kernel(x: torch.Tensor) -> torch.Tensor: out = torch.empty_like(x) ops.my_kernel_forward(out, x) return out ``` ### Pre-flight checklist before declaring a kernel done 1. `grep -rn "pybind11\|PYBIND11\|torch/extension.h\|py::" torch-ext/` returns nothing. 2. `grep -rn "TORCH_LIBRARY(\|TORCH_LIBRARY_FRAGMENT\|PyInit" torch-ext/` returns nothing (only `TORCH_LIBRARY_EXPAND` is allowed). 3. No `setup.py` exists unless generated by `kernel-builder create-pyproject`. 4. `kernel-builder check-config` passes — `[general]` needs a **dash-separated** `name` (never underscores) and a `license`, plus `[torch]` (binding sources) and `[kernel.<name>]` sections. 5. The kernel directory is a git repository with all files committed (Nix refuses non-git builds). 6. The build succeeds: `nix run .#build-and-copy -L`. 7. ABI compliance passes: `kernel-builder check-abi` (after building). ## Quick Start ### Diffusers (Video/Image Generation) **For benchmarking kernel performance:** ```bash # Benchmark with optimized kernels (6% end-to-end speedup) python generate_video.py --use-optimized-kernels # Benchmark baseline with torch.compile (34% speedup) python generate_video.py --no-optimized-kernels --compile # Compare configurations (note: --compile and --use-optimized-kernels are mutually exclusive) python generate_video.py --use-optimized-kernels && \ python generate_video.py --no-optimized-kernels --compile ``` **For a minimal diffusers integration example (~150 lines):** ```bash python scripts/ltx_kernel_injection_example.py ``` ### Transformers (LLMs) **For a minimal transformers integration example (~120 lines):** ```bash python scripts/transformers_injection_example.py ``` ### HuggingFace Kernels Hub **Load pre-compiled kernels from HuggingFace Hub (no local compilation):** ```python from kernels import get_kernel # Load optimized activation kernels activation = get_kernel("kernels-community/activation", version=1) # Use the kernel y = torch.empty_like(x) activation.gelu_fast(y, x) ``` **For a complete HuggingFace Kernels example:** ```bash python scripts/huggingface_kernels_example.py ``` ### Isolated Kernel Micro-benchmarks ```bash python benchmark_rmsnorm.py ``` ## Supported Libraries & Models | Library | Supported Models | Key Kernels | |---------|------------------|-------------| | **diffusers** | LTX-Video, Stable Diffusion, FLUX, DiT | RMSNorm, GEGLU, RoPE, AdaLN | | **transformers** | LLaMA, Mistral, Qwen, Falcon | RMSNorm, Attention | | GPU | Compute Capability | Guide | |-----|-------------------|-------| | H100 | sm_90 | [h100-optimization-guide.md](references/h100-optimization-guide.md) | | A100 | sm_80 | [a100-optimization-guide.md](references/a100-optimization-guide.md) | | T4 | sm_75 | [t4-optimization-guide.md](references/t4-optimization-guide.md) | ## When This Skill Applies Use this skill when: - **Benchmarking kernel performance** against baseline implementations - Writing new CUDA kernels for diffusion models or LLMs - Optimizing existing kernels for H100, A100, or T4 architecture - Implementing custom attention, normalization, or activation layers - Integrating kernels with **diffusers** pipelines (LTX-Video, Stable Diffusion, FLUX, DiT) - Integrating kernels with **transformers** models (LLaMA, Mistral, Qwen) - Debugging kernel performance issues on NVIDIA GPUs ## Working Example Complete working examples ship with the kernels repo under `examples/kernels/` (also at [github.com/huggingface/kernels](https://github.com/huggingface/kernels/tree/main/examples/kernels)): - `relu/` — the canonical minimal kernel: build.toml, flake.nix, `TORCH_LIBRARY_EXPAND` bindings, Python API, `layers/`, tests - `relu-backprop-compile/` — backward pass + `torch.compile` support (fake-op registration) - `silu-and-mul/` — activation kernel following the same layout ## Benchmarking Kernels Use the benchmark script to measure kernel performance: ```bash # Full benchmark with all options python scripts/benchmark_example.py \ --use-optimized-kernels \ --compile \ --batch-size 1 \ --num-frames 161 \ --height 512 \ --width 768 \ --steps 50 \ --warmup-iterations 2 ``` ### Benchmark Script Options | Option | Default | Description | |--------|---------|-------------| | `--use-optimized-kernels` | auto | Use custom H100 CUDA kernels | | `--no-optimized-kernels` | - | Use baseline implementation | | `--compile` | false | Enable torch.compile on transformer | | `--batch-size` | 1 | Number of videos per prompt | | `--num-frames` | 161 | Number of frames to generate | | `--height` | 512 | Video height in pixels | | `--width` | 768 | Video width in pixels | | `--steps` | 50 | Denoising steps | | `--warmup-iterations` | 2 | Warmup runs before benchmark | ### Example Benchmark Results **End-to-End Video Generation (49 frames, 30 steps, H100 80GB):** | Configuration | Time (s) | it/s | Speedup | Notes | |:---|:---:|:---:|:---:|:---| | Baseline (no compile) | 2.87 | 12.58 | 1.00x | Reference | | **Optimized Kernels** | 2.70 | 13.52 | **1.06x** | 6% faster | | Baseline + torch.compile | 2.14 | 19.05 | 1.34x | 34% faster | **Important:** `--use-optimized-kernels` and `--compile` are currently mutually exclusive. Custom kernels require PyTorch custom op registration to work with torch.compile. **Key metrics to capture:** - **Device:** GPU model (e.g., NVIDIA H100 80GB HBM3) - **Precision:** Data type used (e.g., bfloat16) - **Resolution:** Width x Height (e.g., 768x512) - **Frames:** Number of frames generated (e.g., 49, 161) ### RMSNorm Micro-benchmarks The vectorized RMSNorm kernel achieves **2.67x average speedup** over PyTorch baseline: | Shape | Custom (ms) | PyTorch (ms) | Speedup | |:---|:---:|:---:|:---:| | [1×1024×2048] | 0.019 | 0.065 | **3.37x** | | [2×1024×2048] | 0.024 | 0.073 | **3.04x** | | [4×1024×2048] | 0.036 | 0.093 | **2.58x** | | [2×4096×3072] | 0.087 | 0.208 | **2.41x** | | [4×4096×3072] | 0.157 | 0.392 | **2.49x** | **Bandwidth efficiency:** 38% of H100's theoretical 3.35 TB/s **Why end-to-end speedup is smaller:** RMSNorm accounts for ~5% of total compute in LTX-Video. The remaining time is spent in attention (Flash Attention/SDPA), linear projections, and VAE decode. ## Project Structure ``` .claude/skills/cuda-kernels/ ├── scripts/ │ ├── benchmark_example.py # End-to-end video generation benchmark │ ├── benchmark_rmsnorm.py # Isolated RMSNorm micro-benchmark │ ├── ltx_kernel_injection_example.py # Minimal diffusers integration (~150 lines) │ ├── transformers_injection_example.py # Minimal transformers integration (~120 lines) │ └── huggingface_kernels_example.py # HuggingFace Kernels Hub integration ├── references/ │ ├── diffusers-integration.md # Complete diffusers integration guide │ ├── transformers-integration.md # Complete transformers integration guide │ ├── huggingface-kernels-integration.md # HuggingFace Kernels Hub (get_kernel) guide │ ├── troubleshooting.md # Common issues and solutions │ ├── kernel-templates.md # CUDA kernel templates (includes vectorized) │ ├── h100-optimization-guide.md # H100 (Hopper) optimization deep dive │ ├── a100-optimization-guide.md # A100 (Ampere) optimization deep dive │ └── t4-optimization-guide.md # T4 (Turing) optimization deep dive └── SKILL.md # This file examples/kernels/relu/ # Canonical working example (kernels repo) ├── build.toml # kernel-builder build configuration ├── flake.nix # Nix build entry point ├── CARD.md # Kernel card template (becomes README.md) ├── relu_cuda/relu.cu # CUDA kernel source ├── torch-ext/ │ ├── torch_binding.h / .cpp # TORCH_LIBRARY_EXPAND bindings │ └── relu/__init__.py # Python API (+ optional layers/) └── tests/test_relu.py # Kernel tests (nix run .#ci-test) ```