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

#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`:

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:

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

python scripts/ltx_kernel_injection_example.py

Transformers (LLMs)

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

python scripts/transformers_injection_example.py

HuggingFace Kernels Hub

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

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:

python scripts/huggingface_kernels_example.py

Isolated Kernel Micro-benchmarks

python benchmark_rmsnorm.py

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:

# 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

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

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

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

// Load 2 bfloat16 elements at once (32-bit load)
const __nv_bfloat162* vec_input = reinterpret_cast<const __nv_bfloat162*>(row_input);

#pragma unroll 4
for (int i = tid; i < vec_hidden; i += stride) {
    __nv_bfloat162 v = vec_input[i];
    float v0 = __bfloat162float(v.x);
    float v1 = __bfloat162float(v.y);
    sum_sq += v0 * v0 + v1 * v1;
}

FP16 vectorization using `__half2`:

const __half2* vec_input = reinterpret_cast<const __half2*>(row_input);
__half2 v = vec_input[i];
float v0 = __half2float(v.x);
float v1 = __half2float(v.y);

FP32 vectorization using `float4`:

const float4* vec_input = reinterpret_cast<const float4*>(row_input);
float4 v = vec_input[i];
sum_sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;

Warp Shuffle Reductions

template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}

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

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

Editable install for local 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:

kernel-builder create-pyproject -f
pip install wheel
pip install --no-build-isolation -e .

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 = 1
license = "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").

Library Integration

HuggingFace Kernels Hub (get_kernel)

See [huggingface-kernels-integration.md](references/huggingface-kernels-integration.md) for the complete guide.

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.

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](references/diffusers-integration.md) for the complete guide.

Transformers Integration (LLMs)

See [transformers-integration.md](references/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:

from transformers import AutoModelForCausalLM
from ltx_kernels import rmsnorm

def patch_rmsnorm(model):
    for name, module in model.named_modules():
        if 'RMSNorm' in type(module).__name__:
            eps = getattr(module, 'variance_epsilon', None) or getattr(module, 'eps', 1e-6)
            def make_forward(mod, epsilon):
                def forward(x):
                    return rmsnorm(x, mod.weight, eps=epsilon)
                return forward
            module.forward = make_forward(module, eps)

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype=torch.bfloat16)
patch_rmsnorm(model)

Diffusers Critical Pitfalls

1. RMSNorm Weight May Be None

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

# Transformer blocks: NO WEIGHT
self.norm1 = RMSNorm(dim, elementwise_affine=False)

# Attention modules: HAS WEIGHT
self.norm_q = torch.nn.RMSNorm(..., elementwise_affine=True)

Solution: Handle both cases:

has_weight = hasattr(module, 'weight') and module.weight is not None
if has_weight:
    output = rmsnorm(x, module.weight, eps=eps)
else:
    weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    output = rmsnorm(x, weight, eps=eps)

2. Diffusers RMSNorm != torch.nn.RMSNorm

# WRONG - misses diffusers RMSNorm
if isinstance(module, torch.nn.RMSNorm):

# CORRECT - catches all RMSNorm variants
if type(module).__name__ == '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

pipe = LTXPipeline.from_pretrained(...)
pipe.to("cuda")
inject_optimized_kernels(pipe)  # BEFORE offloading
pipe.enable_model_cpu_offload()  # Now safe

Minimal Integration Pattern

from diffusers import LTXPipeline
from ltx_kernels import rmsnorm

def patch_rmsnorm_modules(model):
    """Patch all RMSNorm modules to use custom kernel."""
    for name, module in model.named_modules():
        if type(module).__name__ == 'RMSNorm':
            eps = getattr(module, 'eps', 1e-6)
            has_weight = hasattr(module, 'weight') and module.weight is not None

            if has_weight:
                def make_forward(mod, epsilon):
                    def forward(x):
                        return rmsnorm(x, mod.weight, eps=epsilon)
                    return forward
                module.forward = make_forward(module, eps)
            else:
                def make_forward(epsilon):
                    def forward(x):
                        w = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
                        return rmsnorm(x, w, eps=epsilon)
                    return forward
                module.forward = make_forward(eps)

# Usage
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")
patch_rmsnorm_modules(pipe.transformer)
pipe.enable_model_cpu_offload()

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

# NVIDIA Nsight Systems
nsys profile -o profile python your_script.py

# NVIDIA Nsight Compute
ncu --set full -o metrics python your_script.py

Common Issues

See [troubleshooting.md](references/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:

torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped

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:

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):
    return None  # out-variant op: no shape changes

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

A100 GPU Optimization Guide for Diffusers/Transformers Kernels

Deep dive into A100-specific optimizations for diffusion model and LLM CUDA kernels.

A100 Ampere Architecture Overview

Key Specifications

A100 vs H100 Comparison

Key A100 Features

1. Third-Gen Tensor Cores - FP16, BF16, TF32, INT8, INT4 2. Multi-Instance GPU (MIG) - Partition into up to 7 instances 3. Structural Sparsity - 2:4 sparsity support in tensor cores 4. TF32 Mode - FP32-like range with FP16-like throughput 5. Asynchronous Copy - Overlap compute and memory

Memory Hierarchy Optimization

Global Memory Access Patterns

Same principles as H100, but lower bandwidth makes coalescing even more critical:

// GOOD: Coalesced access
int idx = blockIdx.x * blockDim.x + threadIdx.x;
float val = input[idx];

// BAD: Strided access (even worse on A100 due to lower bandwidth)
int idx = threadIdx.x * stride;
float val = input[idx];

A100 Transaction sizes:

• 32 bytes minimum
• 128 bytes optimal (full warp, FP32)
• Memory-bound kernels more limited by 2.0 TB/s (vs 3.35 TB/s on H100)

Vectorized Memory Access

Same vectorization patterns work on A100:

BFloat16 vectorization:

const __nv_bfloat162* vec_input = reinterpret_cast<const __nv_bfloat162*>(row_input);

#pragma unroll 4
for (int i = tid; i < hidden_size / 2; i += stride) {
    __nv_bfloat162 v = vec_input[i];
    float v0 = __bfloat162float(v.x);
    float v1 = __bfloat162float(v.y);
}

Expected A100 Performance (RMSNorm):

Bandwidth achieved: Target 30-40% of A100's 2.0 TB/s theoretical

L2 Cache Utilization

A100's 40MB L2 cache is still significant:

// For attention: Same block size tuning works
// BLOCK_SIZE_M = 128  (Q block)
// BLOCK_SIZE_N = 64   (K,V block)
// Tiles fit in L2 for reuse

Shared Memory Configuration

A100 supports configurable shared memory per SM:

• 48 KB shared + 80 KB L1 (default)
• 96 KB shared + 32 KB L1
• 164 KB shared + 0 KB L1 (max)

For attention kernels:

// Request max shared memory
cudaFuncSetAttribute(
    attention_forward_kernel,
    cudaFuncAttributeMaxDynamicSharedMemorySize,
    164 * 1024  // 164 KB max on A100
);

Warp-Level Optimizations

Shuffle Instructions

Same warp shuffle patterns work on A100:

template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}

Occupancy Tuning

Block Size Selection for A100

Grid Sizing

For A100 with 108 SMs:

// Aim for multiples of 108 blocks
int num_blocks = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
// Round up to multiple of 108 for full SM utilization
num_blocks = ((num_blocks + 107) / 108) * 108;

Precision and Tensor Cores

TF32 Mode (A100 Specific)

TF32 provides FP32-like range with better throughput:

# Enable TF32 for matmuls (PyTorch)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

// CUDA: Use TF32 tensor core operations
// Automatically enabled for FP32 inputs on Ampere+ with cuBLAS/cuDNN

BF16 vs FP16 on A100

FP16: Good precision, risk of overflow
BF16: Same range as FP32, preferred for training
TF32: Best throughput for FP32-like accuracy (A100 specific)

Build Configuration

build.toml for A100

[general]
name = "ltx-kernels"   # dash-separated; underscores are rejected
backends = ["cuda"]
version = 1
license = "Apache-2.0"

[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"]
cuda-capabilities = ["8.0"]  # sm_80 for A100

Multi-GPU Support (A100 + H100)

[kernel.your_kernel]
backend = "cuda"
src = ["kernel_src/your_kernel.cu"]
cuda-capabilities = ["8.0", "9.0"]  # Both A100 and H100

CUDA Compilation Flags

# For A100 specifically
nvcc -arch=sm_80 -O3 your_kernel.cu

# For both A100 and H100
nvcc -gencode=arch=compute_80,code=sm_80 \
     -gencode=arch=compute_90,code=sm_90 \
     -O3 your_kernel.cu

A100-Specific Optimizations

Async Memory Copy

A100 introduced async memory copy (cp.async):

// Async copy from global to shared memory
__pipeline_memcpy_async(shared_ptr, global_ptr, size);
__pipeline_commit();
__pipeline_wait_prior(0);

Structural Sparsity

A100 tensor cores support 2:4 sparsity (50% zeros):

# PyTorch sparse semi-structured
from torch.sparse import to_sparse_semi_structured
sparse_weight = to_sparse_semi_structured(dense_weight)

Multi-Instance GPU (MIG)

For inference workloads:

# Create MIG instances
nvidia-smi mig -cgi 9,9,9,9,9,9,9 -C
# Creates 7x 5GB instances on A100 40GB

Performance Profiling

Expected Performance (A100 vs H100)

Nsight Profiling

# Same commands work on A100
nsys profile -o a100_profile python your_script.py
ncu --set full -o a100_metrics.ncu-rep python your_script.py

# Key A100 metrics to watch:
# - sm__throughput.avg.pct_of_peak_sustained_elapsed
# - dram__throughput.avg.pct_of_peak_sustained_elapsed
# - gpu__time_duration.avg (kernel time)

Migration from H100 to A100

Code Changes Required

1. Shared Memory: Reduce max shared memory from 192KB to 164KB 2. Grid Size: Adjust for 108 SMs instead of 132 3. No TMA: Can't use Tensor Memory Accelerator 4. No FP8: Must use FP16/BF16 instead

Backward Compatible Pattern

// Works on both A100 and H100
#if __CUDA_ARCH__ >= 900
    // H100-specific optimizations (TMA, etc.)
#else
    // A100/older GPU fallback
#endif

Best Practices Summary (A100)

1. Memory Access: Even more critical due to lower bandwidth 2. Vectorization: Use __nv_bfloat162, __half2, float4 3. TF32: Enable for FP32 workloads for ~8x speedup 4. Block Size: 256 threads is good default 5. Shared Memory: Max 164 KB/SM 6. Grid Size: Multiples of 108 for full utilization 7. Profile: Compare achieved vs theoretical bandwidth 8. Multi-arch: Build for both sm_80 and sm_90

Working Example

cd <your-kernel-project>

# Leave cuda-capabilities unspecified in build.toml unless the kernel
# truly requires specific architectures (A100 is `cuda-capabilities = ["8.0"]`).
nix run .#build-and-copy -L  # Build kernels with kernel-builder

# Run the kernel's test suite
nix run .#ci-test

H100 CUDA Kernels for Diffusers

This skill provides patterns and guidance for developing optimized CUDA kernels targeting NVIDIA H100 GPUs (compute capability 9.0) for use with the HuggingFace diffusers library.

Quick Start

For benchmarking kernel performance:

# 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 isolated kernel micro-benchmarks:

python benchmark_rmsnorm.py

For a minimal integration example (~150 lines):

python scripts/ltx_kernel_injection_example.py

When This Skill Applies

Use this skill when:

• Benchmarking kernel performance against baseline implementations
• Writing new CUDA kernels for diffusion models
• Optimizing existing kernels for H100 architecture
• Implementing custom attention, normalization, or activation layers
• Integrating kernels with diffusers pipelines (LTX-Video, Stable Diffusion, FLUX, DiT)
• Debugging kernel performance issues on H100

Working Example

Complete working examples ship with the kernels repo under examples/kernels/ (e.g. relu/, relu-backprop-compile/). They demonstrate:

• Custom CUDA kernels with the canonical project layout
• Build system setup with build.toml and flake.nix
• PyTorch C++ bindings (TORCH_LIBRARY_EXPAND) and Python API
• Kernel tests runnable via nix run .#ci-test

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

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

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 integration example
├── references/
│   ├── diffusers-integration.md    # Complete integration guide
│   ├── troubleshooting.md          # Common issues and solutions
│   ├── kernel-templates.md         # CUDA kernel templates (includes vectorized)
│   └── h100-optimization-guide.md  # H100 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
├── relu_cuda/relu.cu               # CUDA kernel source
├── torch-ext/                      # TORCH_LIBRARY_EXPAND bindings + Python API
└── tests/                          # Kernel tests

H100 Architecture Reference

Core Kernel Patterns

Vectorized Memory Access (Critical for Performance)

BFloat16 vectorization using `__nv_bfloat162`:

// Load 2 bfloat16 elements at once (32-bit load)
const __nv_bfloat162* vec_input = reinterpret_cast<const __nv_bfloat162*>(row_input);

#pragma unroll 4
for (int i = tid; i < vec_hidden; i += stride) {
    __nv_bfloat162 v = vec_input[i];
    float v0 = __bfloat162float(v.x);
    float v1 = __bfloat162float(v.y);
    sum_sq += v0 * v0 + v1 * v1;
}

FP16 vectorization using `__half2`:

const __half2* vec_input = reinterpret_cast<const __half2*>(row_input);
__half2 v = vec_input[i];
float v0 = __half2float(v.x);
float v1 = __half2float(v.y);

FP32 vectorization using `float4`:

const float4* vec_input = reinterpret_cast<const float4*>(row_input);
float4 v = vec_input[i];
sum_sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;

Warp Shuffle Reductions

template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}

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

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

With Nix (Recommended)

nix run .#build-and-copy --max-jobs 2 --cores 8 -L

Editable install for local 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:

kernel-builder create-pyproject -f
pip install wheel
pip install --no-build-isolation -e .

build.toml Configuration

[general]
# Dash-separated lowercase name (underscores are rejected); license required.
name = "ltx-kernels"
backends = ["cuda"]
version = 1
license = "Apache-2.0"

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

Diffusers Integration

See [diffusers-integration.md](references/diffusers-integration.md) for the complete guide.

Critical Pitfalls

1. RMSNorm Weight May Be None

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

# Transformer blocks: NO WEIGHT
self.norm1 = RMSNorm(dim, elementwise_affine=False)

# Attention modules: HAS WEIGHT
self.norm_q = torch.nn.RMSNorm(..., elementwise_affine=True)

Solution: Handle both cases:

has_weight = hasattr(module, 'weight') and module.weight is not None
if has_weight:
    output = rmsnorm(x, module.weight, eps=eps)
else:
    weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
    output = rmsnorm(x, weight, eps=eps)

2. Diffusers RMSNorm != torch.nn.RMSNorm

# WRONG - misses diffusers RMSNorm
if isinstance(module, torch.nn.RMSNorm):

# CORRECT - catches all RMSNorm variants
if type(module).__name__ == '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

pipe = LTXPipeline.from_pretrained(...)
pipe.to("cuda")
inject_optimized_kernels(pipe)  # BEFORE offloading
pipe.enable_model_cpu_offload()  # Now safe

Minimal Integration Pattern

from diffusers import LTXPipeline
from ltx_kernels import rmsnorm

def patch_rmsnorm_modules(model):
    """Patch all RMSNorm modules to use custom kernel."""
    for name, module in model.named_modules():
        if type(module).__name__ == 'RMSNorm':
            eps = getattr(module, 'eps', 1e-6)
            has_weight = hasattr(module, 'weight') and module.weight is not None

            if has_weight:
                def make_forward(mod, epsilon):
                    def forward(x):
                        return rmsnorm(x, mod.weight, eps=epsilon)
                    return forward
                module.forward = make_forward(module, eps)
            else:
                def make_forward(epsilon):
                    def forward(x):
                        w = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
                        return rmsnorm(x, w, eps=epsilon)
                    return forward
                module.forward = make_forward(eps)

# Usage
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")
patch_rmsnorm_modules(pipe.transformer)
pipe.enable_model_cpu_offload()

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

# NVIDIA Nsight Systems
nsys profile -o profile python your_script.py

# NVIDIA Nsight Compute
ncu --set full -o metrics python your_script.py

Common Issues

See [troubleshooting.md](references/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:

torch._dynamo.exc.Unsupported: Attempted to call function marked as skipped

Workaround options: 1. Use --use-optimized-kernels without --compile (6% speedup) 2. Use --compile without custom kernels (34% speedup) 3. Register kernel as custom op (advanced, requires torch.library)

To register as custom op (for torch.compile compatibility):

import torch

@torch.library.custom_op("ltx_kernels::rmsnorm", mutates_args={"out"})
def rmsnorm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, eps: float) -> None:
    ops.rmsnorm_forward(out, input.contiguous(), weight.contiguous(), eps)

@rmsnorm.register_fake
def _(out, input, weight, eps):
    pass  # No shape changes

See Also

• benchmark_example.py - Benchmarking script for comparing optimized vs baseline - START HERE
• ltx_kernel_injection_example.py - Minimal integration example (~150 lines)
• diffusers-integration.md - Complete integration guide
• troubleshooting.md - Common issues and solutions
• kernel-templates.md - Complete kernel templates
• h100-optimization-guide.md - Deep dive on H100 optimizations
• examples/kernels/relu/ - Canonical working kernel example

Diffusers Pipeline Integration Guide

Complete guide for integrating custom CUDA kernels into HuggingFace diffusers pipelines.

Quick Start: See ltx_kernel_injection_example.py for a minimal working example (~150 lines).

Overview

Diffusers pipelines (LTX-Video, Stable Diffusion, FLUX, DiT) have specific architecture patterns. Understanding these patterns is critical for successful kernel integration.

Model Architecture Analysis

Before integrating kernels, analyze the target model:

# 1. Check pipeline components
from diffusers import LTXPipeline
import inspect

pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
print("Components:", [k for k in dir(pipe) if not k.startswith('_') and hasattr(getattr(pipe, k), 'named_modules')])
# Output: ['transformer', 'vae', 'text_encoder']

# 2. Analyze transformer structure
for name, module in pipe.transformer.named_modules():
    class_name = type(module).__name__
    if 'Norm' in class_name or 'Attention' in class_name or 'GELU' in class_name:
        has_weight = hasattr(module, 'weight') and module.weight is not None
        print(f"{name}: {class_name} (has_weight={has_weight})")

LTX-Video Architecture

Key Components

Kernel Applicability

Integration Pattern

Step 1: Create Optimized Attention Processor

from typing import Optional, Tuple
import torch
from ltx_kernels import rmsnorm

class OptimizedLTXVideoAttnProcessor:
    """
    Custom attention processor using optimized CUDA kernels.

    Replaces RMSNorm operations for Q/K normalization with custom kernel.
    """

    def __call__(
        self,
        attn,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ) -> torch.Tensor:
        # Import here to avoid circular imports
        from diffusers.models.transformers.transformer_ltx import apply_rotary_emb
        from diffusers.models.attention_dispatch import dispatch_attention_fn

        batch_size, sequence_length, _ = (
            hidden_states.shape if encoder_hidden_states is None
            else encoder_hidden_states.shape
        )

        if attention_mask is not None:
            attention_mask = attn.prepare_attention_mask(
                attention_mask, sequence_length, batch_size
            )
            attention_mask = attention_mask.view(
                batch_size, attn.heads, -1, attention_mask.shape[-1]
            )

        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states

        # Q, K, V projections
        query = attn.to_q(hidden_states)
        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)

        # Custom RMSNorm kernel for Q/K normalization
        # NOTE: attn.norm_q and attn.norm_k HAVE weights (elementwise_affine=True)
        query = rmsnorm(query, attn.norm_q.weight, eps=attn.norm_q.eps)
        key = rmsnorm(key, attn.norm_k.weight, eps=attn.norm_k.eps)

        # Apply rotary embeddings (computed by diffusers)
        if image_rotary_emb is not None:
            query = apply_rotary_emb(query, image_rotary_emb)
            key = apply_rotary_emb(key, image_rotary_emb)

        # Reshape for multi-head attention
        query = query.unflatten(2, (attn.heads, -1))
        key = key.unflatten(2, (attn.heads, -1))
        value = value.unflatten(2, (attn.heads, -1))

        # Dispatch attention (PyTorch SDPA or other backends)
        hidden_states = dispatch_attention_fn(
            query, key, value,
            attn_mask=attention_mask,
            dropout_p=0.0,
            is_causal=False,
        )
        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)

        # Output projection
        hidden_states = attn.to_out[0](hidden_states)
        hidden_states = attn.to_out[1](hidden_states)
        return hidden_states

Step 2: Create Module Patcher

import torch.nn as nn
from ltx_kernels import rmsnorm

def patch_rmsnorm_modules(model: nn.Module) -> int:
    """
    Patch all RMSNorm modules to use custom CUDA kernel.

    Handles both:
    - Modules WITH weight (elementwise_affine=True)
    - Modules WITHOUT weight (elementwise_affine=False)

    Returns:
        Number of modules patched.
    """
    patched_count = 0

    for name, module in model.named_modules():
        # Check by class name (not isinstance) to catch diffusers RMSNorm
        if type(module).__name__ == 'RMSNorm':
            eps = getattr(module, 'eps', 1e-6)
            has_weight = hasattr(module, 'weight') and module.weight is not None

            if has_weight:
                def make_patched_forward_with_weight(mod, epsilon):
                    def patched_forward(x):
                        return rmsnorm(x, mod.weight, eps=epsilon)
                    return patched_forward
                module.forward = make_patched_forward_with_weight(module, eps)
            else:
                # No weight (elementwise_affine=False) - use ones
                def make_patched_forward_no_weight(epsilon):
                    def patched_forward(x):
                        weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
                        return rmsnorm(x, weight, eps=epsilon)
                    return patched_forward
                module.forward = make_patched_forward_no_weight(eps)

            patched_count += 1

    return patched_count

Step 3: Create Injection Function

def inject_optimized_kernels(pipe) -> dict:
    """
    Inject custom CUDA kernels into the pipeline.

    Call this AFTER loading model to CUDA, BEFORE enabling CPU offloading.

    Returns:
        dict with counts of patched modules.
    """
    stats = {
        'attention_processors': 0,
        'rmsnorm_modules': 0,
    }

    if not hasattr(pipe, 'transformer'):
        print("WARNING: Pipeline has no 'transformer' attribute!")
        return stats

    transformer = pipe.transformer

    # 1. Replace attention processors
    for name, module in transformer.named_modules():
        if hasattr(module, 'set_processor') and hasattr(module, 'processor'):
            module.set_processor(OptimizedLTXVideoAttnProcessor())
            stats['attention_processors'] += 1

    # 2. Patch RMSNorm modules
    stats['rmsnorm_modules'] = patch_rmsnorm_modules(transformer)

    return stats

Step 4: Use in Script

import torch
from diffusers import LTXPipeline
from diffusers.utils import export_to_video

# Import kernels
from ltx_kernels import rmsnorm

# Load pipeline
pipe = LTXPipeline.from_pretrained(
    "Lightricks/LTX-Video",
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

# Inject kernels BEFORE CPU offloading
stats = inject_optimized_kernels(pipe)
print(f"Attention processors replaced: {stats['attention_processors']}")
print(f"RMSNorm modules patched: {stats['rmsnorm_modules']}")

# Now enable memory optimization
pipe.enable_model_cpu_offload()

# Generate video
output = pipe(
    prompt="A cat sleeping in the sun",
    num_frames=25,
    height=480,
    width=704,
    num_inference_steps=30,
)

export_to_video(output.frames[0], "output.mp4", fps=24)

Model-Specific Differences

LTX-Video

• Uses GELU (not GEGLU)
• RMSNorm in blocks has no weight
• RMSNorm in attention has weight
• Custom 3D RoPE computed by diffusers

Stable Diffusion 3

• Uses GEGLU in FeedForward
• May have different normalization patterns
• Check before assuming kernel applicability

FLUX

• Uses GEGLU
• Different attention patterns
• Verify architecture before patching

Verification

Verify Injection Worked

# Check attention processors
for name, module in pipe.transformer.named_modules():
    if hasattr(module, 'processor'):
        print(f"{name}: {type(module.processor).__name__}")
        break
# Should show: OptimizedLTXVideoAttnProcessor

# Test a forward pass
with torch.inference_mode():
    x = torch.randn(1, 100, 2048, device='cuda', dtype=torch.bfloat16)
    for name, module in pipe.transformer.named_modules():
        if type(module).__name__ == 'RMSNorm':
            out = module(x)
            print(f"RMSNorm forward pass: {x.shape} -> {out.shape}")
            break

Run Full Inference Test

.venv/bin/python generate_video.py --num-frames 9 --steps 5
# Quick test with minimal frames/steps

Troubleshooting

See SKILL.md "Common Issues and Solutions" for:

• Weight is None errors
• isinstance() not working
• GEGLU not being called
• CPU offloading issues

Complete Example

For a self-contained, runnable example that demonstrates all patterns above:

cd <your-kernel-project>     # kernel-builder project with your kernels
nix run .#build-and-copy -L  # Build kernels with kernel-builder
python path/to/skills/cuda-kernels/scripts/ltx_kernel_injection_example.py

This example: 1. Loads LTX-Video pipeline 2. Injects custom kernels 3. Verifies injection worked 4. Generates a short test video

H100 GPU Optimization Guide for Diffusers Kernels

Deep dive into H100-specific optimizations for diffusion model CUDA kernels.

H100 Hopper Architecture Overview

Key Specifications

New Hopper Features

1. Thread Block Clusters - Groups of thread blocks that can cooperate 2. Distributed Shared Memory - Access shared memory across blocks in cluster 3. Tensor Memory Accelerator (TMA) - Hardware-accelerated bulk memory operations 4. FP8 Support - Native 8-bit floating point in tensor cores 5. Asynchronous Execution - More overlap between compute and memory

Memory Hierarchy Optimization

Global Memory Access Patterns

// GOOD: Coalesced access (threads access consecutive addresses)
// Each thread reads 4 bytes, warp reads 128 bytes (one transaction)
int idx = blockIdx.x * blockDim.x + threadIdx.x;
float val = input[idx];

// BAD: Strided access (each thread in warp accesses different cache lines)
int idx = threadIdx.x * stride;  // Avoid stride > 1
float val = input[idx];

Transaction sizes:

• 32 bytes minimum
• 128 bytes optimal (full warp, FP32)
• Align to 128-byte boundaries when possible

Vectorized Memory Access (Critical for Bandwidth)

Vectorized loads/stores dramatically improve memory bandwidth utilization:

BFloat16 vectorization (2x elements per load):

// Load 2 bfloat16 elements at once (32-bit transaction)
const __nv_bfloat162* vec_input = reinterpret_cast<const __nv_bfloat162*>(row_input);

#pragma unroll 4
for (int i = tid; i < hidden_size / 2; i += stride) {
    __nv_bfloat162 v = vec_input[i];
    float v0 = __bfloat162float(v.x);
    float v1 = __bfloat162float(v.y);
    // Process v0, v1...
}

// Write back vectorized
__nv_bfloat162* vec_output = reinterpret_cast<__nv_bfloat162*>(row_output);
__nv_bfloat162 result;
result.x = __float2bfloat16(val0);
result.y = __float2bfloat16(val1);
vec_output[i] = result;

FP16 vectorization:

const __half2* vec_input = reinterpret_cast<const __half2*>(row_input);
__half2 v = vec_input[i];
float v0 = __half2float(v.x);
float v1 = __half2float(v.y);

FP32 vectorization (4x elements per load):

const float4* vec_input = reinterpret_cast<const float4*>(row_input);
float4 v = vec_input[i];
// v.x, v.y, v.z, v.w are 4 consecutive floats

Benchmark Results (RMSNorm on H100):

Bandwidth achieved: 38% of H100's theoretical 3.35 TB/s

L2 Cache Utilization

H100's 50MB L2 cache is significant for diffusion models:

// For attention: Process Q blocks to maximize K,V cache reuse
// K,V tiles stay in L2 while Q block iterates

// Block size tuning for L2:
// BLOCK_SIZE_M = 128  (Q block)
// BLOCK_SIZE_N = 64   (K,V block)
// With head_dim=64, each tile = 128*64*2 = 16KB (FP16)
// Multiple tiles fit in L2 for reuse

Shared Memory Configuration

H100 supports configurable shared memory per SM:

• 96 KB shared + 128 KB L1
• 144 KB shared + 80 KB L1
• 192 KB shared + 32 KB L1

For attention kernels with large tiles:

// Request max shared memory
cudaFuncSetAttribute(
    attention_forward_kernel,
    cudaFuncAttributeMaxDynamicSharedMemorySize,
    192 * 1024  // 192 KB
);

Bank Conflicts

Shared memory has 32 banks (4 bytes per bank):

// Bank conflict example (all threads hit same bank)
__shared__ float data[1024];
float val = data[threadIdx.x * 32];  // BAD: 32-stride = same bank

// No bank conflict
float val = data[threadIdx.x];  // GOOD: consecutive access

// Bank conflict avoidance with padding
__shared__ float data[32][33];  // 33 instead of 32
float val = data[threadIdx.y][threadIdx.x];  // Different banks

Warp-Level Optimizations

Shuffle Instructions

Fastest way to share data within a warp:

// Reduction using shuffles (no shared memory needed)
template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}

// Broadcast from lane 0
float broadcast = __shfl_sync(0xffffffff, val, 0);

// Butterfly shuffle for max
float max_val = __shfl_xor_sync(0xffffffff, val, 16);
val = max(val, max_val);
// ... repeat for 8, 4, 2, 1

Warp-Level Matrix Operations

For small matrices, keep data in registers:

// Example: 4x4 matrix per warp (1 element per thread in first 16 lanes)
// More efficient than shared memory for small sizes

Register Optimization

Register Pressure

H100 allows 255 registers per thread. Monitor usage:

nvcc --ptxas-options=-v your_kernel.cu
# Shows: "Used X registers, Y bytes smem"

Register Tiling

For attention, keep partial results in registers:

// Each thread maintains its own row_max and row_sum
float row_max = -INFINITY;
float row_sum = 0.0f;

// And output accumulator (fits in registers if head_dim is small)
float out_acc[HEAD_DIM];  // Works for head_dim <= ~64

Occupancy Tuning

Calculating Occupancy

Occupancy = Active Warps per SM / Max Warps per SM (64)

Limiting factors:
1. Registers: 65536 registers / (threads_per_block * regs_per_thread)
2. Shared Memory: 192KB / smem_per_block
3. Threads: 2048 / threads_per_block

Block Size Selection

For H100 diffusers kernels:

Occupancy Calculator Usage

# Use CUDA occupancy API
from numba import cuda
import numba.cuda as nb_cuda

@cuda.jit
def my_kernel(...):
    pass

# Get suggested block size
max_block_size = my_kernel.suggest_cooperative_groups_max_block_size()
occupancy = my_kernel.occupancy(max_block_size)

Precision and Numerical Stability

BF16 vs FP16

For diffusion models:

FP16: 1 sign + 5 exponent + 10 mantissa
- Better precision (10 bits)
- Smaller range (±65504)
- Risk of overflow in attention scores

BF16: 1 sign + 8 exponent + 7 mantissa
- Same range as FP32
- Less precision (7 bits)
- Safer for attention (no overflow)
- Preferred for training

Online Softmax for Attention

Numerically stable softmax without materializing full attention matrix:

// Traditional (bad for memory)
// scores = Q @ K^T  // [seq, seq] - huge!
// softmax(scores)
// output = scores @ V

// Online softmax (good)
float row_max = -INFINITY;
float row_sum = 0.0f;

for each K block:
    compute local_scores
    local_max = max(local_scores)

    // Update running statistics
    new_max = max(row_max, local_max)
    rescale = exp(row_max - new_max)

    row_sum = row_sum * rescale + sum(exp(local_scores - new_max))
    row_max = new_max

    // Update output accumulator with rescaling
    out_acc = out_acc * rescale + softmax_scores @ V_block

Mixed Precision Pattern

Use FP32 for reductions, low precision for memory:

// Input in FP16/BF16
float sum = 0.0f;  // Accumulate in FP32
for (int i = tid; i < hidden_size; i += blockDim.x) {
    float val = float(input[i]);  // Cast to FP32
    sum += val * val;
}
// Reduction in FP32
sum = block_reduce_sum(sum);

// Output in FP16/BF16
output[i] = scalar_t(result);  // Cast back

Diffusers-Specific Optimizations

LTX-Video Attention Pattern

LTX-Video uses 3D positional encoding for video:

// Sequence layout: [batch, num_frames * height * width, heads, head_dim]
// Position encoding splits head_dim into temporal + spatial components

// Efficient 3D position decoding
int t_idx = seq_idx / (height * width);
int hw_idx = seq_idx % (height * width);
int h_idx = hw_idx / width;
int w_idx = hw_idx % width;

// Apply different RoPE frequencies to different head_dim ranges
// Typically: head_dim / 3 for each of (t, h, w)

DiT Adaptive LayerNorm

DiT uses timestep-conditioned normalization:

// Formula: norm(x) * weight * (1 + scale) + shift
// scale, shift come from MLP on timestep embedding

// Optimization: Fuse the MLP projection with AdaLN application
// Compute 6 values per block: (scale1, shift1, gate1, scale2, shift2, gate2)
// Apply to attention output and FFN output respectively

GEGLU FFN Pattern

Common in modern transformers:

// Input: [batch, seq, 2*hidden]
// Split into gate and value halves
// Output: gelu(gate) * value

// Memory optimization: Don't materialize intermediate
float gate = float(input[idx]);
float value = float(input[idx + hidden_size]);
float activated = gelu_tanh(gate) * value;
output[idx] = scalar_t(activated);

Profiling and Debugging

NVIDIA Nsight Systems (nsys)

System-wide profiling:

nsys profile -o profile_report python your_script.py

# Key metrics to watch:
# - Kernel duration
# - Memory transfer time
# - GPU idle time
# - Stream utilization

NVIDIA Nsight Compute (ncu)

Detailed kernel analysis:

# Full metrics
ncu --set full -o metrics.ncu-rep python your_script.py

# Specific metrics
ncu --metrics sm__throughput.avg.pct_of_peak_sustained_elapsed,\
dram__throughput.avg.pct_of_peak_sustained_elapsed \
python your_script.py

# Key metrics for diffusers kernels:
# - Achieved occupancy
# - Memory throughput
# - Compute throughput
# - Warp stall reasons

Common Performance Issues

1. Low occupancy: Too many registers or shared memory

• Solution: Reduce register usage, use smaller block sizes

2. Memory bound: Low compute/memory ratio

• Solution: Fuse operations, increase arithmetic intensity

3. Bank conflicts: Shared memory access pattern issues

• Solution: Add padding, change access pattern

4. Warp divergence: Conditional branches within warp

• Solution: Restructure to process similar elements together

5. Launch overhead: Too many small kernels

• Solution: Fuse kernels, use CUDA graphs

CUDA Compilation Flags

# For H100 specifically
nvcc -arch=sm_90 -O3 your_kernel.cu

# Useful flags:
# -maxrregcount=N    Limit registers per thread
# --ptxas-options=-v Print register/smem usage
# -lineinfo          Add debug line info
# --use_fast_math    Fast but less precise math
# -Xptxas -dlcm=ca   Cache global loads in L1

Best Practices Summary

1. Memory Access: Always coalesce, align to 128 bytes 2. Shared Memory: Use for data reuse, watch bank conflicts 3. Registers: Prefer for small, thread-private data 4. Reductions: Use warp shuffles, avoid atomics when possible 5. Precision: BF16 for training, FP16 for inference, FP32 for accumulation 6. Block Size: Start with 256, tune based on occupancy 7. Profile: Use ncu to identify bottlenecks before optimizing 8. Fuse: Combine operations to reduce memory traffic 9. Type Conversions: Always use explicit to_float()/from_float() helpers (PyTorch disables implicit FP16/BF16 conversions)

Working Example

For complete, working kernel projects, see examples/kernels/ in the kernels repo (e.g. relu/, relu-backprop-compile/).

Benchmark results on H100 (custom LTX-Video kernels):

Build and test:

cd <your-kernel-project>
nix run .#build-and-copy -L  # Build kernels with kernel-builder
nix run .#ci-test            # Run the kernel's test suite

HuggingFace Kernels Integration Guide

Complete guide for using and publishing CUDA kernels with the HuggingFace Kernels library (get_kernel).

Quick Start: See huggingface_kernels_example.py for a minimal working example.

Overview

The HuggingFace Kernels library enables dynamic loading of pre-compiled CUDA kernels from the Hugging Face Hub. This eliminates the need for local compilation and ensures compatibility across different Python, PyTorch, and CUDA versions.

Key Benefits:

• No local compilation - Download pre-built binaries
• Version management - Load specific kernel versions
• Multi-version support - Multiple versions coexist in one Python process
• Automatic compatibility - Matches your PyTorch/CUDA configuration

Installation

pip install kernels torch numpy

Requirements:

• PyTorch >= 2.5
• CUDA-capable GPU
• Python 3.8+

Core API

get_kernel

Download and load a kernel from the Hub:

from kernels import get_kernel

# Load a specific major version (the standard way). A bare
# get_kernel(repo_id) raises ValueError — version= or revision= is required.
kernel = get_kernel("kernels-community/activation", version=1)

# Or pin an explicit revision (branch/tag/commit). This is for exceptional cases, using `version` is strongly recommended.
kernel = get_kernel("kernels-community/flash-attn", revision="v2.0.0")

Parameters:

Returns: ModuleType - the imported kernel module

has_kernel

Check if a kernel build exists for your environment:

from kernels import has_kernel

if has_kernel("kernels-community/activation", version=1):
    kernel = get_kernel("kernels-community/activation", version=1)
else:
    print("No compatible build available")

get_local_kernel

Load a locally built kernel (useful for development — no Hub access, no version needed). Pass the kernel project root; it resolves variants from <path> or <path>/build:

from pathlib import Path
from kernels import get_local_kernel

# Load the freshly built kernel (after `kernel-builder build-and-copy -L`)
kernel = get_local_kernel(Path("/path/to/my-kernel"))

Alternatively, the LOCAL_KERNELS environment variable redirects get_kernel() itself to a local build — production integration code can then be tested unchanged:

LOCAL_KERNELS="my-username/my-kernel=/path/to/my-kernel" python app.py
# get_kernel("my-username/my-kernel") now loads the local build/
# (the override skips the Hub, trust checks, and version resolution)

load_kernel & get_locked_kernel

For reproducible, offline-capable deployments using lockfiles:

from kernels import load_kernel, get_locked_kernel

# Load using a lockfile
kernel = load_kernel("lockfile.json")

# Get kernel with lock
kernel = get_locked_kernel("kernels-community/activation", lockfile="kernel.lock")

Usage Examples

1. Basic Activation Kernel

import torch
from kernels import get_kernel

# Load activation kernels from Hub
activation = get_kernel("kernels-community/activation", version=1)

# Create test tensor
x = torch.randn((10, 10), dtype=torch.float16, device="cuda")

# Execute kernel (output tensor must be pre-allocated)
y = torch.empty_like(x)
activation.gelu_fast(y, x)

print(y)

2. Flash Attention

import torch
from kernels import get_kernel

flash_attn = get_kernel("kernels-community/flash-attn", version=1)

# Check available functions
print(dir(flash_attn))

# Usage depends on specific kernel API

3. RMSNorm Kernel

import torch
from kernels import get_kernel

layer_norm = get_kernel("kernels-community/triton-layer-norm", version=1)

# Apply RMSNorm
x = torch.randn(2, 1024, 2048, dtype=torch.bfloat16, device="cuda")
weight = torch.ones(2048, dtype=torch.bfloat16, device="cuda")
out = layer_norm.rms_norm(x, weight, eps=1e-6)

4. Integration with Transformers Models

import torch
import torch.nn as nn
from kernels import get_kernel

# Load RMSNorm kernel
rmsnorm_kernel = get_kernel("kernels-community/triton-layer-norm", version=1)

def patch_rmsnorm_with_hub_kernel(model):
    """Patch model's RMSNorm to use Hub kernel."""
    for name, module in model.named_modules():
        if 'RMSNorm' in type(module).__name__:
            eps = getattr(module, 'variance_epsilon', None) or getattr(module, 'eps', 1e-6)

            def make_forward(mod, epsilon):
                def forward(hidden_states):
                    return rmsnorm_kernel.rms_norm(hidden_states, mod.weight, eps=epsilon)
                return forward

            module.forward = make_forward(module, eps)

5. Integration with Diffusers Pipelines

import torch
from diffusers import LTXPipeline
from kernels import get_kernel, has_kernel

# Load kernel if available
if has_kernel("kernels-community/activation", version=1):
    activation = get_kernel("kernels-community/activation", version=1)

    def patch_activations(model):
        # Patch GELU activations with optimized kernel
        for name, module in model.named_modules():
            if isinstance(module, torch.nn.GELU):
                def make_forward():
                    def forward(x):
                        out = torch.empty_like(x)
                        activation.gelu_fast(out, x)
                        return out
                    return forward
                module.forward = make_forward()

# Use with pipeline
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)
pipe.to("cuda")
patch_activations(pipe.transformer)

Publishing Kernels to Hub

Project Structure

Scaffold new projects with kernel-builder init — it generates this layout (plus benchmarks/, example.py) with a valid build.toml and an initialized git repository:

kernel-builder init --name my-username/my-kernel

my-kernel/
├── build.toml           # Build configuration
├── flake.nix            # Required: kernel-builder's Nix build entry point
├── CARD.md              # Kernel card template (uploaded as README.md)
├── my_kernel_cuda/
│   └── my_kernel.cu     # CUDA source (any dir name; listed in build.toml src)
├── torch-ext/
│   ├── torch_binding.cpp
│   ├── torch_binding.h
│   └── my_kernel/
│       └── __init__.py
└── tests/
    └── test_my_kernel.py

build.toml Configuration

[general]
# Dash-separated lowercase name (underscores are rejected by check-config);
# the Python package dir is torch-ext/<name with dashes -> underscores>.
name = "my-kernel"
backends = ["cuda"]
version = 1
license = "Apache-2.0"

[general.hub]
# Hub repository to upload to (used by `kernel-builder build-and-upload`);
# together with `version` this selects the version branch (e.g. v1).
repo-id = "my-username/my-kernel"

[torch]
src = [
  "torch-ext/torch_binding.cpp",
  "torch-ext/torch_binding.h"
]

[kernel.my_kernel]
backend = "cuda"
src = ["my_kernel_cuda/my_kernel.cu"]
depends = ["torch"]

# Leave cuda-capabilities unspecified unless the kernel truly requires
# specific architectures — over-constraining produces non-compliant builds.
# cuda-capabilities = ["9.0", "10.0", "12.0"]

Torch Bindings

This is the only supported binding pattern. Do NOT use pybind11 (PYBIND11_MODULE, #include <torch/extension.h>) or TORCH_LIBRARY with a hardcoded namespace — both fail under kernel-builder's ABI3 build. See "Hard Constraints" in SKILL.md.

torch_binding.h:

#pragma once
#include <torch/torch.h>

void my_kernel_forward(torch::Tensor &out, torch::Tensor const &input);

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

Python Wrapper

torch-ext/my_kernel/__init__.py:

from typing import Optional
import torch
from ._ops import ops

def forward(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Apply my custom kernel."""
    if out is None:
        out = torch.empty_like(x)
    ops.my_kernel_forward(out, x)
    return out

Layers (Optional)

A kernel can also export torch.nn.Module layers that kernels.kernelize() swaps in for matching model layers. Per the kernel requirements, layers must be pure: subclass torch.nn.Module, define no constructor, no class variables (except has_backward / can_torch_compile), and no methods other than forward. Put them in torch-ext/my_kernel/layers.py and export the module from the main __init__.py:

# torch-ext/my_kernel/layers.py
import torch
from ._ops import ops

class MyKernelLayer(torch.nn.Module):
    has_backward: bool = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = torch.empty_like(x)
        ops.my_kernel_forward(out, x)
        return out

A layer's forward may also use member variables (e.g. weight, bias) that are defined by the layer it extends. Since the layer defines no constructor, these are not assigned here — but their types can be annotated as class-level type hints purely for type checking:

# torch-ext/my_kernel/layers.py
import torch
from ._ops import ops

class RMSNorm(torch.nn.Module):
    has_backward: bool = True
    can_torch_compile: bool = True

    # Defined by the layer being extended; annotated for type checking.
    weight: torch.Tensor
    variance_epsilon: float

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # `rms_norm` is defined by the kernel.
        return ops.rms_norm(
            hidden_states,
            self.weight,
            self.variance_epsilon,
        )

# torch-ext/my_kernel/__init__.py
from . import layers

__all__ = [..., "layers"]

Building and Publishing

Using kernel-builder (Nix):

# Build the kernel locally (run inside the kernel directory)
kernel-builder build-and-copy -L

# Build and upload to the Hub in one go
kernel-builder build-and-upload

The target repository is determined by the repo-id and version fields in build.toml (see above). Uploads go to a `kernel`-type Hub repository (the first-class kernel repository type), not a model repo — the owning user or org needs kernel-creation access, requested via huggingface.co/settings/account ("Request Kernels Creation"). If a CARD.md template is present in the source repo, it is filled and uploaded as the README.md.

Editable install for local development (never hand-write a setup.py — torch.utils.cpp_extension/pybind11 cannot build under ABI3):

kernel-builder create-pyproject -f
pip install wheel
pip install --no-build-isolation -e .

Available Community Kernels

Popular kernels from kernels-community:

Browse all kernels: https://huggingface.co/kernels-community

Inspecting Kernel Functions

Kernel function signatures vary by implementation. Always inspect before use:

from kernels import get_kernel

kernel = get_kernel("kernels-community/activation", version=1)

# List available functions
print(dir(kernel))
# ['gelu_fast', 'gelu_new', 'silu', ...]

# Check function signature (if available)
import inspect
print(inspect.signature(kernel.gelu_fast))

Caching and Offline Usage

Downloaded kernels are cached in the HuggingFace Hub cache directory:

• Default: ~/.cache/huggingface/hub/
• Override: Set HF_HOME environment variable

For offline usage:

import os
os.environ["HF_HUB_OFFLINE"] = "1"

# Will only use cached kernels
kernel = get_kernel("kernels-community/activation", version=1)

Best Practices

1. Check availability first:

   if has_kernel("kernels-community/my-kernel", version=1):
       kernel = get_kernel("kernels-community/my-kernel", version=1)
   else:
       # Fallback to PyTorch implementation

2. Always pass `version=` (it is required, not optional):

   kernel = get_kernel("kernels-community/activation", version=1)

Version branches (v1, v2, ...) never break the kernel API, so pinning the major version keeps code working while still receiving fixes.

3. Use lockfiles for production:

   kernel = load_kernel("kernel.lock")

4. Pre-allocate output tensors:

   # Most kernels require pre-allocated outputs
   out = torch.empty_like(x)
   kernel.function(out, x)

5. Test with your exact environment:

   # Print environment info
   import torch
   print(f"PyTorch: {torch.__version__}")
   print(f"CUDA: {torch.version.cuda}")
   print(f"GPU: {torch.cuda.get_device_name()}")

Troubleshooting

No compatible build found

from kernels import has_kernel, get_kernel

if not has_kernel("kernels-community/my-kernel", version=1):
    print("No build for your PyTorch/CUDA version")
    print(f"PyTorch: {torch.__version__}, CUDA: {torch.version.cuda}")

Import errors after loading

# Always inspect available functions
kernel = get_kernel("kernels-community/activation", version=1)
print(dir(kernel))  # Check what's actually available

Version conflicts

# Explicitly specify version
kernel_v1 = get_kernel("repo/kernel", version=1)
kernel_v2 = get_kernel("repo/kernel", version=2)
# Both can coexist in the same process

See Also

• HuggingFace Kernels Documentation
• HuggingFace Kernels GitHub
• Kernel Builder Documentation
• Community Kernels
• Blog: Learn the Kernel Hub in 5 Minutes

CUDA Kernel Templates for H100 Diffusers

Complete, copy-paste ready templates for implementing new kernels.

Hard constraints (ABI3 / kernel-builder): Never use pybind11 — that includes #include <torch/extension.h>, which pulls in pybind11 transitively, and PYBIND11_MODULE. Never hand-write a setup.py with torch.utils.cpp_extension. Never use TORCH_LIBRARY(...) with a hardcoded namespace. Bindings MUST use TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) + REGISTER_EXTENSION (Template 4), built via build.toml. See "Hard Constraints" in SKILL.md for the full disallowed list.

CRITICAL: Type Conversion Helpers

PyTorch compiles with `-D__CUDA_NO_HALF_OPERATORS__` which disables implicit FP16/BF16 conversions. You MUST include these helpers in every kernel file:

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>

// Type conversion helpers - REQUIRED for PyTorch compatibility
__device__ __forceinline__ float to_float(float x) { return x; }
__device__ __forceinline__ float to_float(__half x) { return __half2float(x); }
__device__ __forceinline__ float to_float(__nv_bfloat16 x) { return __bfloat162float(x); }

__device__ __forceinline__ float from_float(float x, float*) { return x; }
__device__ __forceinline__ __half from_float(float x, __half*) { return __float2half(x); }
__device__ __forceinline__ __nv_bfloat16 from_float(float x, __nv_bfloat16*) { return __float2bfloat16(x); }

Usage in kernels:

// Read with conversion
float val = to_float(input[idx]);

// Write with conversion (use nullptr cast for type deduction)
output[idx] = from_float(result, (scalar_t*)nullptr);

Template 1: Element-wise Operation (RoPE style)

Use this pattern for operations that process elements independently.

/*
 * Element-wise kernel template for H100 (sm_90)
 */

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cmath>

constexpr int BLOCK_SIZE = 256;

// Type conversion helpers (include in every .cu file)
__device__ __forceinline__ float to_float(float x) { return x; }
__device__ __forceinline__ float to_float(__half x) { return __half2float(x); }
__device__ __forceinline__ float to_float(__nv_bfloat16 x) { return __bfloat162float(x); }

__device__ __forceinline__ float from_float(float x, float*) { return x; }
__device__ __forceinline__ __half from_float(float x, __half*) { return __float2half(x); }
__device__ __forceinline__ __nv_bfloat16 from_float(float x, __nv_bfloat16*) { return __float2bfloat16(x); }

template <typename scalar_t>
__global__ void your_elementwise_kernel(
    scalar_t* __restrict__ output,
    const scalar_t* __restrict__ input,
    const int total_elements
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < total_elements) {
        float val = to_float(input[idx]);

        // Your computation here
        float result = val;  // Replace with actual operation

        output[idx] = from_float(result, (scalar_t*)nullptr);
    }
}

// C++ entry points
extern "C" {

void your_kernel_forward_fp16(
    __half* output,
    const __half* input,
    int total_elements,
    cudaStream_t stream
) {
    const int num_blocks = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
    your_elementwise_kernel<__half><<<num_blocks, BLOCK_SIZE, 0, stream>>>(
        output, input, total_elements
    );
}

void your_kernel_forward_bf16(
    __nv_bfloat16* output,
    const __nv_bfloat16* input,
    int total_elements,
    cudaStream_t stream
) {
    const int num_blocks = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
    your_elementwise_kernel<__nv_bfloat16><<<num_blocks, BLOCK_SIZE, 0, stream>>>(
        output, input, total_elements
    );
}

void your_kernel_forward_fp32(
    float* output,
    const float* input,
    int total_elements,
    cudaStream_t stream
) {
    const int num_blocks = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
    your_elementwise_kernel<float><<<num_blocks, BLOCK_SIZE, 0, stream>>>(
        output, input, total_elements
    );
}

}

Template 2: Row-wise Reduction (LayerNorm style)

Use for operations requiring reduction across a dimension (normalization, softmax).

Basic Version (Scalar Loads)

/*
 * Row-wise reduction kernel template for H100 (sm_90)
 */

#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cmath>

constexpr int WARP_SIZE = 32;
constexpr int MAX_THREADS = 1024;

// Type conversion helpers
__device__ __forceinline__ float to_float(float x) { return x; }
__device__ __forceinline__ float to_float(__half x) { return __half2float(x); }
__device__ __forceinline__ float to_float(__nv_bfloat16 x) { return __bfloat162float(x); }

__device__ __forceinline__ float from_float(float x, float*) { return x; }
__device__ __forceinline__ __half from_float(float x, __half*) { return __float2half(x); }
__device__ __forceinline__ __nv_bfloat16 from_float(float x, __nv_bfloat16*) { return __float2bfloat16(x); }

template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}

template <typename T>
__device__ __forceinline__ T block_reduce_sum(T val) {
    __shared__ T shared[32];
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;

    val = warp_reduce_sum(val);

    if (lane == 0) shared[wid] = val;
    __syncthreads();

    val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : T(0);
    if (wid == 0) val = warp_reduce_sum(val);

    return val;
}

template <typename scalar_t>
__global__ void your_reduction_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ weight,
    scalar_t* __restrict__ output,
    const int hidden_size,
    const float eps
) {
    const int row = blockIdx.x;
    const int tid = threadIdx.x;

    const scalar_t* row_input = input + row * hidden_size;
    scalar_t* row_output = output + row * hidden_size;

    // Step 1: Compute reduction (e.g., sum of squares for RMSNorm)
    float sum_sq = 0.0f;
    for (int i = tid; i < hidden_size; i += blockDim.x) {
        float val = to_float(row_input[i]);
        sum_sq += val * val;
    }
    sum_sq = block_reduce_sum(sum_sq);

    // Step 2: Compute normalization factor
    __shared__ float s_factor;
    if (tid == 0) {
        s_factor = rsqrtf(sum_sq / hidden_size + eps);
    }
    __syncthreads();
    float factor = s_factor;

    // Step 3: Apply normalization
    for (int i = tid; i < hidden_size; i += blockDim.x) {
        float normalized = to_float(row_input[i]) * factor;
        row_output[i] = from_float(normalized * to_float(weight[i]), (scalar_t*)nullptr);
    }
}

Optimized Version: Vectorized BF16 RMSNorm (2.67x faster)

/*
 * Vectorized RMSNorm kernel for BF16 - H100 optimized
 * Uses __nv_bfloat162 for 2-element vectorized memory access
 * Achieves 2.67x speedup over scalar version
 */
__global__ void rmsnorm_kernel_bf16_vectorized(
    __nv_bfloat16* __restrict__ output,
    const __nv_bfloat16* __restrict__ input,
    const __nv_bfloat16* __restrict__ weight,
    const int hidden_size,
    const float eps
) {
    extern __shared__ char smem[];
    float* shared = reinterpret_cast<float*>(smem);

    const int row = blockIdx.x;
    const int tid = threadIdx.x;
    const int stride = blockDim.x;

    const __nv_bfloat16* row_input = input + row * hidden_size;
    __nv_bfloat16* row_output = output + row * hidden_size;

    // Phase 1: Compute sum of squares with bf16x2 vectorized loads
    float sum_sq = 0.0f;
    const int vec_hidden = hidden_size / 2;
    const __nv_bfloat162* vec_input = reinterpret_cast<const __nv_bfloat162*>(row_input);

    #pragma unroll 4
    for (int i = tid; i < vec_hidden; i += stride) {
        __nv_bfloat162 v = vec_input[i];
        float v0 = __bfloat162float(v.x);
        float v1 = __bfloat162float(v.y);
        sum_sq += v0 * v0 + v1 * v1;
    }

    // Handle odd element if hidden_size is odd
    if (hidden_size % 2 == 1 && tid == 0) {
        float v = __bfloat162float(row_input[hidden_size - 1]);
        sum_sq += v * v;
    }

    // Reduce across block
    sum_sq = block_reduce_sum(sum_sq, shared);

    // Compute RMS inverse
    __shared__ float rms_inv;
    if (tid == 0) {
        float mean_sq = sum_sq / static_cast<float>(hidden_size);
        rms_inv = rsqrtf(mean_sq + eps);
    }
    __syncthreads();

    const float factor = rms_inv;

    // Phase 2: Apply normalization with bf16x2 vectorized stores
    const __nv_bfloat162* vec_weight = reinterpret_cast<const __nv_bfloat162*>(weight);
    __nv_bfloat162* vec_output = reinterpret_cast<__nv_bfloat162*>(row_output);

    #pragma unroll 4
    for (int i = tid; i < vec_hidden; i += stride) {
        __nv_bfloat162 v_in = vec_input[i];
        __nv_bfloat162 v_w = vec_weight[i];

        float v0 = __bfloat162float(v_in.x);
        float v1 = __bfloat162float(v_in.y);
        float w0 = __bfloat162float(v_w.x);
        float w1 = __bfloat162float(v_w.y);

        __nv_bfloat162 result;
        result.x = __float2bfloat16(v0 * factor * w0);
        result.y = __float2bfloat16(v1 * factor * w1);
        vec_output[i] = result;
    }

    // Handle odd element
    if (hidden_size % 2 == 1 && tid == 0) {
        float v = __bfloat162float(row_input[hidden_size - 1]);
        float w = __bfloat162float(weight[hidden_size - 1]);
        row_output[hidden_size - 1] = __float2bfloat16(v * factor * w);
    }
}

// Launch configuration for vectorized kernel
void rmsnorm_forward_bf16(
    __nv_bfloat16* output,
    const __nv_bfloat16* input,
    const __nv_bfloat16* weight,
    const int batch_size,
    const int seq_len,
    const int hidden_size,
    const float eps,
    cudaStream_t stream
) {
    const int num_rows = batch_size * seq_len;
    // Divide by 2 for vectorized access
    int threads = min(hidden_size / 2, MAX_THREADS);
    threads = max(threads, WARP_SIZE);
    threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;

    size_t smem_size = ((threads + WARP_SIZE - 1) / WARP_SIZE) * sizeof(float);

    if (hidden_size % 2 == 0 && hidden_size >= 64) {
        rmsnorm_kernel_bf16_vectorized<<<num_rows, threads, smem_size, stream>>>(
            output, input, weight, hidden_size, eps
        );
    } else {
        // Fallback to scalar kernel for small/odd sizes
        rmsnorm_kernel<__nv_bfloat16><<<num_rows, threads, smem_size, stream>>>(
            output, input, weight, hidden_size, eps
        );
    }
}

// C++ entry points
extern "C" {

void your_reduction_forward_fp16(
    const __half* input,
    const __half* weight,
    __half* output,
    int batch_size,
    int hidden_size,
    float eps,
    cudaStream_t stream
) {
    int threads = min(hidden_size, MAX_THREADS);
    threads = (threads + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE;

    your_reduction_kernel<__half><<<batch_size, threads, 0, stream>>>(
        input, weight, output, hidden_size, eps
    );
}

void your_reduction_forward_bf16(
    const __nv_bfloat16* input,
    const __nv_bfloat16* weight,
    __nv_bfloat16* output,
    int batch_size,
    int hidden_size,
    float eps,
    cudaStream_t stream
) {
    int threads = min(hidden_size, MAX_THREADS);
    threads = (threads + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE;

    your_reduction_kernel<__nv_bfloat16><<<batch_size, threads, 0, stream>>>(
        input, weight, output, hidden_size, eps
    );
}

void your_reduction_forward_fp32(
    const float* input,
    const float* weight,
    float* output,
    int batch_size,
    int hidden_size,
    float eps,
    cudaStream_t stream
) {
    int threads = min(hidden_size, MAX_THREADS);
    threads = (threads + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE;

    your_reduction_kernel<float><<<batch_size, threads, 0, stream>>>(
        input, weight, output, hidden_size, eps
    );
}

}

Template 3: Tiled Matrix Operation (Attention style)

Use for operations requiring shared memory tiling (matmul, attention).

/*
 * Tiled matrix operation template for H100 (sm_90)
 */

#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cmath>

// Block sizes optimized for H100 L2 cache
constexpr int BLOCK_M = 128;
constexpr int BLOCK_N = 64;
constexpr int BLOCK_K = 64;
constexpr int NUM_WARPS = 8;

template <typename T>
__device__ __forceinline__ T warp_reduce_max(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val = max(val, __shfl_xor_sync(0xffffffff, val, offset));
    }
    return val;
}

template <typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, offset);
    }
    return val;
}

template <typename scalar_t>
__global__ void your_tiled_kernel(
    const scalar_t* __restrict__ A,  // [batch, M, K]
    const scalar_t* __restrict__ B,  // [batch, K, N]
    scalar_t* __restrict__ C,        // [batch, M, N]
    const int batch_size,
    const int M,
    const int N,
    const int K
) {
    // Shared memory for tiles
    extern __shared__ char shared_mem[];
    scalar_t* tile_A = reinterpret_cast<scalar_t*>(shared_mem);
    scalar_t* tile_B = tile_A + BLOCK_M * BLOCK_K;

    const int batch_idx = blockIdx.z;
    const int block_row = blockIdx.y;
    const int block_col = blockIdx.x;

    const int tid = threadIdx.x;

    // Base offsets for this batch
    const scalar_t* batch_A = A + batch_idx * M * K;
    const scalar_t* batch_B = B + batch_idx * K * N;
    scalar_t* batch_C = C + batch_idx * M * N;

    // Initialize accumulator
    float acc[BLOCK_M / (NUM_WARPS * 32)][BLOCK_N / 32] = {0};

    // Iterate over K dimension tiles
    for (int k_tile = 0; k_tile < (K + BLOCK_K - 1) / BLOCK_K; k_tile++) {
        // Cooperative loading of tiles to shared memory
        for (int i = tid; i < BLOCK_M * BLOCK_K; i += blockDim.x) {
            int row = i / BLOCK_K;
            int col = i % BLOCK_K;
            int global_row = block_row * BLOCK_M + row;
            int global_col = k_tile * BLOCK_K + col;

            if (global_row < M && global_col < K) {
                tile_A[i] = batch_A[global_row * K + global_col];
            } else {
                tile_A[i] = scalar_t(0);
            }
        }

        for (int i = tid; i < BLOCK_K * BLOCK_N; i += blockDim.x) {
            int row = i / BLOCK_N;
            int col = i % BLOCK_N;
            int global_row = k_tile * BLOCK_K + row;
            int global_col = block_col * BLOCK_N + col;

            if (global_row < K && global_col < N) {
                tile_B[i] = batch_B[global_row * N + global_col];
            } else {
                tile_B[i] = scalar_t(0);
            }
        }
        __syncthreads();

        // Compute partial results
        // (Simplified - real implementation would use register tiling)
        #pragma unroll
        for (int k = 0; k < BLOCK_K; k++) {
            // Your tiled computation here
        }
        __syncthreads();
    }

    // Write results
    // (Implementation depends on your specific needs)
}

// C++ entry points follow same pattern as above

Template 4: PyTorch Binding

// torch_binding.cpp
// IMPORTANT: Include CUDA headers for __half and __nv_bfloat16 types
// NEVER include <torch/extension.h> — it pulls in pybind11, which cannot
// build under the limited API (ABI3) that kernel-builder requires.

#include <torch/torch.h>
#include <torch/library.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>      // Required for __half
#include <cuda_bf16.h>      // Required for __nv_bfloat16
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "registration.h"   // Generated by kernel-builder — do not write by hand

extern "C" {
void your_kernel_forward_fp16(__half* out, const __half* in, int n, cudaStream_t);
void your_kernel_forward_bf16(__nv_bfloat16* out, const __nv_bfloat16* in, int n, cudaStream_t);
void your_kernel_forward_fp32(float* out, const float* in, int n, cudaStream_t);
}

void your_kernel_forward(
    torch::Tensor& output,
    const torch::Tensor& input
) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(output.is_cuda(), "output must be a CUDA tensor");

    const int total_elements = input.numel();

    const at::cuda::CUDAGuard device_guard(input.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (input.scalar_type() == at::kHalf) {
        your_kernel_forward_fp16(
            input.data_ptr(), output.data_ptr(),
            total_elements, stream
        );
    } else if (input.scalar_type() == at::kBFloat16) {
        your_kernel_forward_bf16(
            input.data_ptr(), output.data_ptr(),
            total_elements, stream
        );
    } else if (input.scalar_type() == at::kFloat) {
        your_kernel_forward_fp32(
            static_cast<const float*>(input.data_ptr()),
            static_cast<float*>(output.data_ptr()),
            total_elements, stream
        );
    } else {
        TORCH_CHECK(false, "Unsupported dtype");
    }
}

// Registration — the ONLY supported pattern. Do NOT use PYBIND11_MODULE
// or TORCH_LIBRARY(my_namespace, ...): kernel-builder suffixes the op
// namespace with a build hash, so only TORCH_EXTENSION_NAME resolves.
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("your_kernel_forward(Tensor! out, Tensor input) -> ()");
  ops.impl("your_kernel_forward", torch::kCUDA, &your_kernel_forward);
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

For kernels with multiple backends in one build.toml, guard each ops.impl with the backend macro kernel-builder defines, so the same binding file compiles for every variant:

#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
  ops.impl("your_kernel_forward", torch::kCUDA, &your_kernel_forward);
#endif

Template 5: Python API

# In torch-ext/ltx_kernels/__init__.py
# Always go through the generated _ops module (relative import).
# NEVER call torch.ops.ltx_kernels.* directly — the real namespace is
# suffixed with a build hash and only _ops knows it.
from typing import Optional

import torch

from ._ops import ops

def your_kernel(
    input: torch.Tensor,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    Your kernel description.

    Args:
        input: Input tensor [batch, seq, hidden]
        out: Optional pre-allocated output tensor

    Returns:
        Output tensor [batch, seq, hidden]
    """
    if out is None:
        out = torch.empty_like(input)

    ops.your_kernel_forward(out, input.contiguous())
    return out

Template 6: build.toml Entry

[kernel.your_kernel]
backend = "cuda"
depends = []
src = ["kernel_src/your_kernel.cu"]
cuda-capabilities = ["9.0"]

Template 7: Test Case

# In tests/test_kernels.py

import torch
import pytest
from ltx_kernels import your_kernel

@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("shape", [(2, 1024, 2048), (1, 4096, 4096)])
def test_your_kernel(dtype, shape):
    device = "cuda"
    input = torch.randn(shape, dtype=dtype, device=device)

    # Reference implementation
    expected = your_reference_implementation(input)

    # Kernel implementation
    output = your_kernel(input)

    # Compare
    rtol = 1e-2 if dtype == torch.float16 else 1e-4
    atol = 1e-3 if dtype == torch.float16 else 1e-5
    torch.testing.assert_close(output, expected, rtol=rtol, atol=atol)

def test_your_kernel_with_preallocated():
    device = "cuda"
    dtype = torch.bfloat16
    shape = (2, 1024, 2048)

    input = torch.randn(shape, dtype=dtype, device=device)
    output = torch.empty_like(input)

    result = your_kernel(input, out=output)

    assert result is output  # Verify in-place

Working Example Reference

For complete, working implementations following these templates, see the examples in the kernels repo (examples/kernels/, also on GitHub):

examples/kernels/relu/
├── build.toml
├── flake.nix
├── CARD.md
├── relu_cuda/relu.cu             # CUDA kernel source
├── torch-ext/
│   ├── torch_binding.h / .cpp    # TORCH_LIBRARY_EXPAND bindings
│   └── relu/__init__.py          # Python API (+ layers/)
└── tests/test_relu.py

relu-backprop-compile/ additionally shows backward-pass registration and torch.compile fake ops. To scaffold a fresh project with this exact layout:

kernel-builder init --name my-username/my-kernel

Build and test with:

cd examples/kernels/relu
nix run .#build-and-copy -L
nix run .#ci-test

For an editable dev install, generate the project files with kernel-builder (never hand-write a setup.py):

kernel-builder create-pyproject -f
pip install wheel && pip install --no-build-isolation -e .

T4 GPU Optimization Guide for Diffusers/Transformers Kernels

Deep dive into T4-specific optimizations for diffusion model and LLM CUDA kernels. The T4 is a Turing architecture GPU commonly found in cloud instances (GCP, AWS, etc.) and is popular for inference workloads.

T4 Turing Architecture Overview

Key Specifications

T4 vs A100 vs H100 Comparison

Key T4 Constraints

1. No BFloat16 - Must use FP16 for half-precision 2. Limited Memory - 16GB requires careful batching 3. Lower Bandwidth - 320 GB/s limits memory-bound kernels 4. Fewer SMs - Less parallelism, smaller grid sizes 5. Smaller Shared Memory - 64 KB/SM limits tile sizes 6. Half Max Threads - 1024/SM instead of 2048

Memory Considerations

FP16 Instead of BF16

Critical: T4 does not support BF16. Use FP16:

# T4: Use FP16
model = model.to(torch.float16)

# A100/H100: Can use BF16
# model = model.to(torch.bfloat16)

// T4: Use __half, NOT __nv_bfloat16
const __half2* vec_input = reinterpret_cast<const __half2*>(row_input);
__half2 v = vec_input[i];
float v0 = __half2float(v.x);
float v1 = __half2float(v.y);

// Write back
__half2 result;
result.x = __float2half(val0);
result.y = __float2half(val1);
vec_output[i] = result;

Memory-Bound Kernel Optimization

T4's 320 GB/s bandwidth is the main bottleneck. Maximize arithmetic intensity:

// BAD: Low arithmetic intensity (memory bound on T4)
for (int i = tid; i < size; i += stride) {
    output[i] = input[i] * scale;  // 1 multiply per 2 loads
}

// BETTER: Fuse operations to increase arithmetic intensity
for (int i = tid; i < size; i += stride) {
    float val = input[i];
    val = val * scale + bias;
    val = max(val, 0.0f);  // ReLU
    output[i] = val;  // More ops per memory access
}

Vectorized Memory Access

Even more critical on T4 due to lower bandwidth:

FP16 vectorization (2x elements per load):

const __half2* vec_input = reinterpret_cast<const __half2*>(row_input);

#pragma unroll 4
for (int i = tid; i < hidden_size / 2; i += stride) {
    __half2 v = vec_input[i];
    float v0 = __half2float(v.x);
    float v1 = __half2float(v.y);
    sum_sq += v0 * v0 + v1 * v1;
}

FP32 vectorization (4x elements per load):

const float4* vec_input = reinterpret_cast<const float4*>(row_input);
float4 v = vec_input[i];
// v.x, v.y, v.z, v.w are 4 consecutive floats

Expected T4 Performance

Bandwidth achieved: Target 40-50% of T4's 320 GB/s theoretical

Shared Memory Configuration

T4 supports configurable shared memory per SM:

• 32 KB shared + 32 KB L1
• 48 KB shared + 16 KB L1
• 64 KB shared + 0 KB L1 (max)

For T4, use smaller tile sizes:

// Request shared memory (max 64 KB on T4)
cudaFuncSetAttribute(
    kernel,
    cudaFuncAttributeMaxDynamicSharedMemorySize,
    64 * 1024  // 64 KB max on T4
);

Tile Size Adjustments

Reduce tile sizes compared to A100/H100:

// H100/A100 attention tile sizes
// BLOCK_SIZE_M = 128, BLOCK_SIZE_N = 64

// T4 attention tile sizes (smaller due to shared memory limits)
constexpr int BLOCK_SIZE_M = 64;   // Reduced
constexpr int BLOCK_SIZE_N = 32;   // Reduced
constexpr int BLOCK_SIZE_K = 32;   // Reduced

Occupancy Tuning

Block Size Selection for T4

Due to max 1024 threads/SM (vs 2048 on A100/H100):

Grid Sizing

For T4 with 40 SMs:

// Aim for multiples of 40 blocks
int num_blocks = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
// Round up to multiple of 40 for full SM utilization
num_blocks = ((num_blocks + 39) / 40) * 40;

Numerical Stability with FP16

FP16 has smaller range than BF16, requiring more care:

FP16: 1 sign + 5 exponent + 10 mantissa
- Range: ±65504
- Risk of overflow in attention scores!

BF16: 1 sign + 8 exponent + 7 mantissa (NOT AVAILABLE ON T4)
- Range: Same as FP32

Attention Score Scaling

// Scale attention scores to prevent FP16 overflow
float scale = 1.0f / sqrtf((float)head_dim);
// For T4 FP16: May need additional scaling
// scale *= 0.125f;  // Extra scaling if overflow occurs

Mixed Precision Pattern

Always accumulate in FP32:

// Input in FP16 (T4)
float sum = 0.0f;  // Accumulate in FP32
for (int i = tid; i < hidden_size; i += blockDim.x) {
    float val = __half2float(input[i]);  // Convert to FP32
    sum += val * val;
}
// Reduction in FP32
sum = block_reduce_sum(sum);

// Output in FP16
output[i] = __float2half(result);

Build Configuration

build.toml for T4

[general]
name = "ltx-kernels"   # dash-separated; underscores are rejected
backends = ["cuda"]
version = 1
license = "Apache-2.0"

[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"]
cuda-capabilities = ["7.5"]  # sm_75 for T4

Multi-GPU Support (T4 + A100 + H100)

[kernel.your_kernel]
backend = "cuda"
src = ["kernel_src/your_kernel.cu"]
cuda-capabilities = ["7.5", "8.0", "9.0"]  # T4, A100, H100

CUDA Compilation Flags

# For T4 specifically
nvcc -arch=sm_75 -O3 your_kernel.cu

# For T4 + A100 + H100
nvcc -gencode=arch=compute_75,code=sm_75 \
     -gencode=arch=compute_80,code=sm_80 \
     -gencode=arch=compute_90,code=sm_90 \
     -O3 your_kernel.cu

T4-Specific Optimizations

INT8 Quantization

T4 tensor cores support INT8 for fast inference:

# PyTorch dynamic quantization
from torch.quantization import quantize_dynamic
model_int8 = quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

TensorRT Optimization

T4 is commonly used with TensorRT:

import torch_tensorrt

# Compile model for T4
trt_model = torch_tensorrt.compile(
    model,
    inputs=[torch_tensorrt.Input(shape=[1, 3, 224, 224], dtype=torch.float16)],
    enabled_precisions={torch.float16}
)

Batch Size Considerations

With only 16GB memory:

# Calculate max batch size
# Model: ~3GB (FP16)
# Activations: ~2GB per batch
# Max batch size: (16 - 3) / 2 ≈ 6

# Use gradient checkpointing for training
model.gradient_checkpointing_enable()

Memory Management

16GB Memory Strategies

1. Enable CPU Offload

pipe.enable_model_cpu_offload()

2. Use Sequential Processing

pipe.enable_sequential_cpu_offload()

3. Reduce Resolution/Frames

# Lower resolution for T4
output = pipe(
    prompt="...",
    height=256,   # Reduced from 512
    width=384,    # Reduced from 768
    num_frames=9  # Reduced from 49
)

4. FP16 Everywhere

pipe = pipeline.from_pretrained(model_id, torch_dtype=torch.float16)

Performance Profiling

Nsight Profiling

nsys profile -o t4_profile python your_script.py
ncu --set full -o t4_metrics.ncu-rep python your_script.py

# Key T4 metrics:
# - Memory throughput (target 40-50% of 320 GB/s)
# - SM utilization (target high with 40 SMs)
# - Occupancy (max 1024 threads/SM)

Common T4 Bottlenecks

1. Memory Bandwidth - 320 GB/s is the main limit 2. Limited Memory - 16GB requires careful management 3. No BF16 - Must handle FP16 overflow risks 4. Smaller Tiles - 64KB shared memory limits

Migration from H100/A100 to T4

Required Changes

1. Precision: BF16 → FP16 2. Shared Memory: Reduce tile sizes (192→64 KB) 3. Grid Size: Adjust for 40 SMs 4. Occupancy: Account for 1024 max threads/SM 5. Memory: Handle 16GB limit

Conditional Compilation

#if __CUDA_ARCH__ >= 800
    // A100/H100: Use BF16
    typedef __nv_bfloat16 half_t;
    typedef __nv_bfloat162 half2_t;
#else
    // T4/Turing: Use FP16
    typedef __half half_t;
    typedef __half2 half2_t;
#endif

Runtime Detection

import torch

def get_optimal_config():
    capability = torch.cuda.get_device_capability()

    if capability >= (9, 0):  # H100
        return {"dtype": torch.bfloat16, "batch_size": 8}
    elif capability >= (8, 0):  # A100
        return {"dtype": torch.bfloat16, "batch_size": 4}
    else:  # T4 and older
        return {"dtype": torch.float16, "batch_size": 1}

Best Practices Summary (T4)

1. Use FP16: BF16 not supported, handle overflow carefully 2. Vectorization: Critical due to low bandwidth 3. Smaller Tiles: 64 KB shared memory limit 4. Grid Size: Multiples of 40 for full utilization 5. Block Size: 256 threads is good default 6. Occupancy: Max 1024 threads/SM 7. Memory: Plan for 16GB limit 8. INT8: Consider quantization for inference 9. Profile: Focus on memory throughput

Working Example

cd <your-kernel-project>

# Leave cuda-capabilities unspecified in build.toml unless the kernel
# truly requires specific architectures (T4 is sm_75).
nix run .#build-and-copy -L  # Build kernels with kernel-builder

# Run the kernel's test suite
nix run .#ci-test

T4 Cloud Instance Notes

T4 is optimized for inference, not training. Consider A100/H100 for training workloads.