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 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

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:

#pragma once
#include <torch/torch.h>
 
void my_kernel_forward(torch::Tensor &out, torch::Tensor const &input);

torch-ext/torch_binding.cpp:

torch-ext/my_kernel/__init__.py:

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:

For a minimal diffusers integration example (~150 lines):

Transformers (LLMs)

For a minimal transformers integration example (~120 lines):

HuggingFace Kernels Hub

Load pre-compiled kernels from HuggingFace Hub (no local compilation):

For a complete HuggingFace Kernels example:

Isolated Kernel Micro-benchmarks

Supported Libraries & Models

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):

• 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:

Benchmark Script Options

Example Benchmark Results

End-to-End Video Generation (49 frames, 30 steps, H100 80GB):

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:

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

GPU Architecture Reference

H100 (Hopper) - Primary Target

Quick Comparison (H100 vs A100 vs T4)

See detailed guides: H100 | A100 | T4

Core Kernel Patterns

Vectorized Memory Access (Critical for Performance)

BFloat16 vectorization using __nv_bfloat162:

FP16 vectorization using __half2:

FP32 vectorization using float4:

Warp Shuffle Reductions

Block Sizes for Attention

• BLOCK_SIZE_M = 128, BLOCK_SIZE_N = 64, BLOCK_SIZE_K = 64
• NUM_WARPS = 8

Thread Configuration

For element-wise ops (RoPE, GEGLU):

For reduction ops (LayerNorm, RMSNorm) with vectorization:

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:

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)

Build and publish to the Hub in one go

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

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").

Library Integration

HuggingFace Kernels Hub (get_kernel)

See huggingface-kernels-integration.md for the complete guide.

Load pre-compiled, optimized kernels directly from HuggingFace Hub without local compilation:

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.

Popular community kernels:

• kernels-community/activation - GELU, SiLU, etc.
• kernels-community/flash-attn - Flash Attention 2
• kernels-community/triton-layer-norm - LayerNorm, RMSNorm

Diffusers Integration (Video/Image Generation)

See diffusers-integration.md for the complete guide.

Transformers Integration (LLMs)

See transformers-integration.md for the complete guide.

Key differences from diffusers:

• Transformers RMSNorm always has weights (no elementwise_affine=False)
• Use 'RMSNorm' in class_name to match LlamaRMSNorm, MistralRMSNorm, etc.
• Check for variance_epsilon (LLaMA) or eps (others) for epsilon
• No set_processor() pattern - use Flash Attention 2 instead

Minimal transformers pattern:

Diffusers Critical Pitfalls

1. RMSNorm Weight May Be None

LTX-Video uses elementwise_affine=False for some RMSNorm modules:

Solution: Handle both cases:

2. Diffusers RMSNorm != torch.nn.RMSNorm

3. LTX-Video Uses GELU, Not GEGLU

LTX-Video uses activation_fn="gelu-approximate". Don't patch GEGLU for LTX-Video.

4. Inject Kernels BEFORE CPU Offloading

Minimal Integration Pattern

Kernel-Specific Guidelines

RMSNorm

• Input layout: [..., hidden_size]
• Epsilon default: 1e-6
• Weight may be None if elementwise_affine=False
• Vectorization: Use __nv_bfloat162 for BF16, __half2 for FP16, float4 for FP32
• Performance: 2.67x faster than PyTorch with vectorized implementation
• Bandwidth: Achieves ~38% of H100's 3.35 TB/s theoretical bandwidth

RoPE

• 1D: [batch, seq, heads, head_dim] - for text
• 3D: [batch, t*h*w, heads, head_dim] - for video
• LTX-Video computes its own RoPE via LTXVideoRotaryPosEmbed

GEGLU vs GELU

• GEGLU: Input [batch, seq, 2*hidden] -> Output [batch, seq, hidden]
• GELU: Standard activation
• LTX-Video uses GELU, NOT GEGLU

AdaLN

• Formula: norm(x) * weight * (1 + scale) + shift
• Used in DiT blocks for conditioning

Performance Profiling

Common Issues

See troubleshooting.md for all common issues and solutions.

Quick fixes:

• "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:

Workaround options:

1. Use --use-optimized-kernels without --compile (6% speedup)
2. Use --compile without custom kernels (34% speedup)
3. 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:

See Also

Scripts

• benchmark_example.py - Benchmarking script for comparing optimized vs baseline - START HERE
• ltx_kernel_injection_example.py - Minimal diffusers integration (~150 lines)
• transformers_injection_example.py - Minimal transformers/LLM integration (~120 lines)
• huggingface_kernels_example.py - HuggingFace Kernels Hub integration

Integration Guides

• huggingface-kernels-integration.md - HuggingFace Kernels Hub (get_kernel) - load pre-compiled kernels
• diffusers-integration.md - Complete diffusers pipeline integration
• transformers-integration.md - Complete transformers/LLM integration

GPU Optimization Guides

• h100-optimization-guide.md - H100 (Hopper, sm_90) deep dive
• a100-optimization-guide.md - A100 (Ampere, sm_80) deep dive
• t4-optimization-guide.md - T4 (Turing, sm_75) deep dive

Reference

• troubleshooting.md - Common issues and solutions
• kernel-templates.md - Complete kernel templates
• examples/kernels/relu/ - Canonical working kernel example (bindings, layers, tests)

External Resources

• HuggingFace Kernels Documentation
• HuggingFace Kernels GitHub
• Community Kernels on Hub