Provides guidance for writing and benchmarking optimized CUDA kernels for NVIDIA GPUs (H100, A100, T4) targeting HuggingFace diffusers and transformers libraries. Kernels must be kernel-builder/ABI3-compliant: no pybind11, no setup.py, TORCH_LIBRARY_EXPAND bindings only. Supports models like LTX-Video, Stable Diffusion, LLaMA, Mistral, and Qwen. Includes integration with HuggingFace Kernels Hub (get_kernel) for loading pre-compiled kernels. Includes benchmarking scripts to compare kernel performance against baseline implementations.
Provides guidance for writing and benchmarking optimized CUDA kernels for NVIDIA GPUs (H100, A100, T4) targeting HuggingFace diffusers and transformers libraries. Kernels must be kernel-builder/ABI3-compliant: no pybind11, no setup.py, TORCH_LIBRARY_EXPAND bindings only. Supports models like LTX-Video, Stable Diffusion, LLaMA, Mistral, and Qwen. Includes integration with HuggingFace Kernels Hub (get_kernel) for loading pre-compiled kernels. Includes benchmarking scripts to compare kernel performance against baseline implementations.
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 and meet the Kernel Hub 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.
grep -rn "TORCH_LIBRARY(\|TORCH_LIBRARY_FRAGMENT\|PyInit" torch-ext/ returns nothing (only TORCH_LIBRARY_EXPAND is allowed).
No setup.py exists unless generated by kernel-builder create-pyproject.
kernel-builder check-config passes — [general] needs a dash-separatedname (never underscores) and a license, plus [torch] (binding sources) and [kernel.<name>] sections.
The kernel directory is a git repository with all files committed (Nix refuses non-git builds).
The build succeeds: nix run .#build-and-copy -L.
ABI compliance passes: kernel-builder check-abi (after building).
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:
# 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.
constexpr int BLOCK_SIZE = 256;
int num_blocks = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
For reduction ops (LayerNorm, RMSNorm) with vectorization:
// Divide by 2 for bf16/fp16 vectorized access
int threads = min(hidden_size / 2, MAX_THREADS);
threads = max(threads, WARP_SIZE);
threads = (threads + 32 - 1) / 32 * 32; // Round to warp boundary
Supported Data Types
All kernels support three precision modes:
__half (FP16) - Default for inference
__nv_bfloat16 (BF16) - Preferred for training
float (FP32) - Reference/debugging
Building Kernels
Scaffold a new kernel project
Start new kernels with kernel-builder init instead of creating files by hand — it generates the compliant layout in one shot:
kernel-builder init --name my-username/my-kernel
This creates build.toml (valid dash-separated name, license, [general.hub] repo-id already wired), flake.nix, torch-ext/ with compilable torch_binding.{h,cpp} and the Python package, a <name>_cuda/ kernel source dir, tests/, benchmarks/, example.py, and CARD.md — and it initializes a git repository (required for builds). Then replace the stub kernel with your own sources and update the src lists in build.toml.
With Nix (Recommended)
nix run .#build-and-copy --max-jobs 2 --cores 8 -L
Build and publish to the Hub in one go
kernel-builder build-and-upload
The target repo is set by repo-id under [general.hub] and version under [general] in build.toml. Uploads go to a kernel-type Hub repository (not a model repo); the owning user/org needs kernel-creation access ("Request Kernels Creation" at huggingface.co/settings/account).
Local build for development
Never hand-write a setup.py (it leads to torch.utils.cpp_extension/pybind11, which cannot build under ABI3). Let kernel-builder generate the project files, then build with setup.py build_kernel (no pip install/editable install needed):
This builds the kernel and puts the output in build, which can be loaded directly with kernels.get_local_kernel(Path("build")). Inside kernel-builder devshell/testshell, LOCAL_KERNELS is set automatically so get_kernel("<repo-id>") resolves to this local build.
build.toml Configuration
[general]# Name MUST be dash-separated lowercase (my-kernel), never underscores —# `kernel-builder check-config` rejects underscores. The Python package# lives at torch-ext/<name with dashes replaced by underscores>.name = "ltx-kernels"backends = ["cuda"]
version = 1license = "Apache-2.0"# required field[general.hub]# Hub repo for `kernel-builder build-and-upload`; with `version` this# selects the version branch (e.g. v1).repo-id = "my-username/ltx-kernels"[torch]src = [
"torch-ext/torch_binding.cpp",
"torch-ext/torch_binding.h"
]
[kernel.your_kernel]backend = "cuda"src = ["kernel_src/your_kernel.cu"]
depends = ["torch"]
# Only constrain cuda-capabilities when the kernel truly requires it —# do not over-specify.
The kernel directory must be a git repository with files committed (git init && git add -A && git commit) — Nix refuses to build non-git kernels ("Kernel is not in a git repository").
Load pre-compiled, optimized kernels directly from HuggingFace Hub without local compilation:
from kernels import get_kernel, has_kernel
# Check availability and load — Hub loads REQUIRE version= (or revision=);# a bare get_kernel(repo_id) raises ValueError.if has_kernel("kernels-community/activation", version=1):
activation = get_kernel("kernels-community/activation", version=1)
# Use the kernel
x = torch.randn((4, 4), dtype=torch.float16, device="cuda")
y = torch.empty_like(x)
activation.gelu_fast(y, x)
Key functions:
get_kernel(repo_id, version=N) - Download and load kernel from Hub; version= (major version) or revision= (branch/tag/commit) is required
has_kernel(repo_id, version=N) - Check if compatible build exists
get_local_kernel(Path("path/to/kernel-project")) - Load a local build (looks in <path> and <path>/build) — use during development
Testing local builds through the get_kernel() code path: set LOCAL_KERNELS="org/name=/path/to/kernel-project" and call get_kernel("org/name") unchanged — the override short-circuits the Hub entirely (no download, no version needed), so integration code can be tested verbatim against a local build.
"NoneType has no attribute contiguous": RMSNorm weight is None, create ones
isinstance() not matching: Use type(module).__name__ instead
GEGLU not called: Model uses GELU, not GEGLU
Patching doesn't persist: Inject before enable_model_cpu_offload()
torch.compile fails with custom kernels: See below
torch.compile Compatibility
Custom CUDA kernels and torch.compile are mutually exclusive unless you register the kernel as a PyTorch custom op.
Error message:
torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped
Workaround options:
Use --use-optimized-kernels without --compile (6% speedup)
Use --compile without custom kernels (34% speedup)
Add a fake/meta implementation for the C++-registered op (see below)
To make the op torch.compile-compatible: ops registered via TORCH_LIBRARY_EXPAND in C++ are already proper custom ops — do NOT re-wrap them with @torch.library.custom_op in Python. Just register a fake (meta) implementation using the generated _ops.py helpers:
import torch
from ._ops import ops, add_op_namespace_prefix
@torch.library.register_fake(add_op_namespace_prefix("rmsnorm_forward"))def_(out, input, weight, eps):
returnNone# out-variant op: no shape changes
See Also
Scripts
benchmark_example.py - Benchmarking script for comparing optimized vs baseline - START HERE
