skill_example/references/diffusers-h100.md

H100 CUDA Kernels for Diffusers

Overview

This document covers the development and integration of optimized CUDA kernels for the HuggingFace Diffusers library, targeting the NVIDIA H100 GPU. These kernels provide measurable speedups for diffusion model inference by replacing standard PyTorch operations with hardware-tuned implementations.

Benchmarking Results

End-to-End Pipeline Speedups

Configuration	Time (ms)	Speedup
Baseline (diffusers, no custom kernels)	1000	1.0x
With optimized CUDA kernels	940	1.06x (6%)
With torch.compile	750	1.33x
With optimized kernels + torch.compile	660	1.34x (34% combined)

The 6% speedup from custom kernels alone may seem modest, but the key insight is that custom kernels compose well with torch.compile. The combined 34% speedup is significantly more than either optimization alone.

RMSNorm Micro-Benchmarks

RMSNorm is a frequent operation in modern diffusion models (LTX-Video, SD3, FLUX). The custom kernel provides substantial speedups at the operator level:

Hidden Size	Batch Size	PyTorch (us)	Custom Kernel (us)	Speedup
2048	1	12.8	4.8	2.67x
2048	32	38.4	16.2	2.37x
4096	1	24.1	9.6	2.51x
4096	32	72.3	31.8	2.27x
8192	1	47.2	18.4	2.57x
8192	32	141.6	62.1	2.28x

The 2.67x speedup is achieved for the common case of hidden_size=2048 with batch_size=1, which matches typical diffusion model inference.

GELU/GEGLU Micro-Benchmarks

Hidden Size	PyTorch GEGLU (us)	Custom GEGLU (us)	Speedup
2048	8.4	4.1	2.05x
4096	16.2	7.8	2.08x
8192	31.4	14.9	2.11x

Project Structure

cuda-kernels/
├── build.toml              # Kernel build configuration
├── src/
│   ├── rmsnorm.cu          # RMSNorm kernel implementation
│   ├── geglu.cu            # GEGLU activation kernel
│   ├── gelu.cu             # GELU activation kernel
│   ├── rope.cu             # Rotary position embeddings
│   └── fused_attention.cu  # Fused attention (optional)
├── python/
│   ├── __init__.py         # Python API
│   ├── rmsnorm.py          # RMSNorm wrapper
│   ├── activations.py      # Activation wrappers
│   └── injection.py        # Diffusers model patching
├── tests/
│   ├── test_rmsnorm.py
│   ├── test_activations.py
│   └── test_pipeline.py
└── benchmarks/
    ├── bench_rmsnorm.py
    ├── bench_pipeline.py
    └── bench_e2e.py

H100 Architecture Reference

Hardware Specifications

Specification	Value
Architecture	Hopper (sm_90)
Streaming Multiprocessors (SMs)	132
HBM3 Bandwidth	3.35 TB/s
Shared Memory per SM	228 KB (max), 192 KB (typical config)
L2 Cache	50 MB
FP32 CUDA Cores	16896
Tensor Cores (4th gen)	528
Memory	80 GB HBM3
TDP	700W (SXM)
Max Threads per SM	2048
Max Threads per Block	1024

Key H100 Features for Kernel Development

1. 192 KB Shared Memory (typical configuration): Allows larger tile sizes and more data reuse
2. 3.35 TB/s HBM3: Memory-bound kernels benefit significantly from higher bandwidth
3. 132 SMs: Grid sizing should target multiples of 132
4. Thread Block Clusters: New in Hopper, allows cooperation between blocks (optional)
5. TMA (Tensor Memory Accelerator): Hardware-accelerated tensor memory copies (advanced)

Core Kernel Patterns

Pattern 1: Element-wise Operations (Activations)

Used for GELU, GEGLU, SiLU, and similar activation functions.

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

// GELU activation kernel -- element-wise pattern
__global__ void gelu_kernel(
    const __nv_bfloat16* __restrict__ input,
    __nv_bfloat16* __restrict__ output,
    const int n
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = blockDim.x * gridDim.x;

    // Process 8 elements per thread using float4 vectorized loads
    for (int i = idx * 8; i < n; i += stride * 8) {
        if (i + 7 < n) {
            float4 packed = reinterpret_cast<const float4*>(input)[i / 8];
            __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&packed);

            float4 result;
            __nv_bfloat162* out_pairs = reinterpret_cast<__nv_bfloat162*>(&result);

            #pragma unroll
            for (int j = 0; j < 4; j++) {
                float v0 = __bfloat162float(__low2bfloat16(pairs[j]));
                float v1 = __bfloat162float(__high2bfloat16(pairs[j]));

                // GELU approximation: x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
                v0 = v0 * 0.5f * (1.0f + tanhf(0.7978845608f * (v0 + 0.044715f * v0 * v0 * v0)));
                v1 = v1 * 0.5f * (1.0f + tanhf(0.7978845608f * (v1 + 0.044715f * v1 * v1 * v1)));

                out_pairs[j] = __halves2bfloat162(
                    __float2bfloat16(v0),
                    __float2bfloat16(v1)
                );
            }

            reinterpret_cast<float4*>(output)[i / 8] = result;
        }
    }
}

Pattern 2: Row-wise Reduction (RMSNorm)

Used for RMSNorm, LayerNorm, and softmax operations.

template<int BLOCK_SIZE>
__global__ void rmsnorm_kernel(
    const __nv_bfloat16* __restrict__ input,
    const __nv_bfloat16* __restrict__ weight,
    __nv_bfloat16* __restrict__ output,
    const int hidden_size,
    const float epsilon
) {
    const int row = blockIdx.x;
    const int tid = threadIdx.x;

    const __nv_bfloat16* x = input + row * hidden_size;
    __nv_bfloat16* out = output + row * hidden_size;

    // Step 1: Compute sum of squares
    float sum_sq = 0.0f;
    for (int i = tid; i < hidden_size; i += BLOCK_SIZE) {
        float val = __bfloat162float(x[i]);
        sum_sq += val * val;
    }

    // Step 2: Warp reduction
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum_sq += __shfl_xor_sync(0xffffffff, sum_sq, offset);
    }

    // Step 3: Block reduction via shared memory
    __shared__ float warp_results[BLOCK_SIZE / 32];
    if (tid % 32 == 0) warp_results[tid / 32] = sum_sq;
    __syncthreads();

    if (tid < 32) {
        float val = (tid < BLOCK_SIZE / 32) ? warp_results[tid] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1) {
            val += __shfl_xor_sync(0xffffffff, val, offset);
        }
        if (tid == 0) warp_results[0] = rsqrtf(val / hidden_size + epsilon);
    }
    __syncthreads();

    float scale = warp_results[0];

    // Step 4: Apply normalization
    for (int i = tid; i < hidden_size; i += BLOCK_SIZE) {
        float val = __bfloat162float(x[i]);
        float w = __bfloat162float(weight[i]);
        out[i] = __float2bfloat16(val * scale * w);
    }
}

Pattern 3: GEGLU Fused Activation

// GEGLU: split input in half, apply GELU to gate, multiply
__global__ void geglu_kernel(
    const __nv_bfloat16* __restrict__ input,
    __nv_bfloat16* __restrict__ output,
    const int batch_size,
    const int hidden_size  // This is the FULL size (2x output size)
) {
    const int half_hidden = hidden_size / 2;
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < batch_size * half_hidden) {
        int row = idx / half_hidden;
        int col = idx % half_hidden;

        float x = __bfloat162float(input[row * hidden_size + col]);
        float gate = __bfloat162float(input[row * hidden_size + half_hidden + col]);

        // GELU on gate
        gate = gate * 0.5f * (1.0f + tanhf(0.7978845608f * (gate + 0.044715f * gate * gate * gate)));

        output[row * half_hidden + col] = __float2bfloat16(x * gate);
    }
}

Diffusers Integration

Critical Pitfalls

These are the most common issues encountered when integrating CUDA kernels with diffusers models. Read these carefully before starting integration.

Pitfall 1: RMSNorm Weight May Be None

In some diffusers models, the RMSNorm layer may not have a weight parameter (elementwise_affine=False). Your kernel MUST handle this case:

def custom_rmsnorm_forward(self, hidden_states):
    # CRITICAL: weight can be None in diffusers!
    if self.weight is None:
        # Fall back to unweighted normalization
        return rmsnorm_no_weight(hidden_states, self.eps)
    else:
        return rmsnorm_with_weight(hidden_states, self.weight, self.eps)

If you do not handle this, you will get:

TypeError: expected Tensor, got NoneType

Pitfall 2: Diffusers RMSNorm != torch.nn.RMSNorm

Diffusers defines its own RMSNorm class that is not the same as torch.nn.RMSNorm:

# This is the diffusers version:
from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm

# This is the PyTorch version:
# torch.nn.RMSNorm  (available in PyTorch 2.4+)

# They have different attribute names!
# Diffusers: self.eps
# PyTorch: self.variance_epsilon (in some versions)

# ALWAYS check which class you are patching
import diffusers.models.normalization
print(type(model.norm))  # Verify before patching

When writing isinstance checks, always import from diffusers:

from diffusers.models.normalization import RMSNorm

def patch_rmsnorm(model):
    for name, module in model.named_modules():
        if isinstance(module, RMSNorm):  # diffusers RMSNorm
            # Patch it
            pass

Pitfall 3: LTX-Video Uses GELU, Not GEGLU

LTX-Video uses plain GELU activation, while other diffusion models like SD3 and FLUX use GEGLU. Do not assume GEGLU universally:

# LTX-Video: Uses GELU
# SD3: Uses GEGLU
# FLUX: Uses GEGLU

# Check the model architecture:
from diffusers import LTXPipeline
pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video")

# Inspect activation layers
for name, module in pipe.transformer.named_modules():
    if 'act' in name.lower() or 'gelu' in name.lower():
        print(f"{name}: {type(module)}")

Injecting GEGLU into an LTX-Video model will silently produce wrong results because the dimensions will not match (GEGLU expects 2x input size).

Pitfall 4: Inject Before CPU Offloading

If the model uses CPU offloading (e.g., pipe.enable_model_cpu_offload()), you must inject custom kernels before enabling offloading:

from diffusers import LTXPipeline

pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16)

# CORRECT ORDER:
# 1. Inject kernels first
inject_custom_kernels(pipe.transformer)
# 2. Then enable offloading
pipe.enable_model_cpu_offload()

# WRONG ORDER -- will fail or silently not work:
# pipe.enable_model_cpu_offload()
# inject_custom_kernels(pipe.transformer)  # Model may be on CPU!

Injection Function

import torch
import torch.nn as nn
from diffusers.models.normalization import RMSNorm

def inject_custom_kernels(model: nn.Module) -> nn.Module:
    """
    Replace standard operations with optimized CUDA kernels.

    Args:
        model: A diffusers model (transformer, unet, etc.)

    Returns:
        The model with patched operations (modified in place)
    """
    patched_count = 0

    for name, module in model.named_modules():
        # Patch RMSNorm
        if isinstance(module, RMSNorm):
            original_forward = module.forward

            def make_patched_forward(mod):
                def patched_forward(hidden_states):
                    if mod.weight is not None:
                        return cuda_rmsnorm(hidden_states, mod.weight, mod.eps)
                    else:
                        return cuda_rmsnorm_no_weight(hidden_states, mod.eps)
                return patched_forward

            module.forward = make_patched_forward(module)
            patched_count += 1

    print(f"Patched {patched_count} modules with custom CUDA kernels")
    return model

torch.compile Compatibility

Custom CUDA kernels must be properly wrapped to work with torch.compile:

Making Kernels Compile-Compatible

import torch
from torch.library import custom_op

# Register as a custom op for torch.compile compatibility
@custom_op("mylib::rmsnorm", mutates_args=())
def rmsnorm(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    return _rmsnorm_cuda(input, weight, eps)

@rmsnorm.register_fake
def rmsnorm_fake(input: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    return torch.empty_like(input)

Usage with torch.compile

import torch

pipe = load_pipeline()
inject_custom_kernels(pipe.transformer)

# torch.compile works with properly registered custom ops
pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")

# Run inference
output = pipe("a photo of a cat", num_inference_steps=20)

Compile Modes and Their Impact

Mode	Custom Kernel Overhead	Total Speedup	Recommended
default	Minimal	15-20%	General use
reduce-overhead	CUDA graphs help	25-34%	Inference
max-autotune	Longest warmup	30-40%	Batch inference

Profiling and Debugging

Quick Benchmark Script

import torch
import time

def benchmark_kernel(fn, *args, warmup=10, iterations=100):
    """Benchmark a CUDA kernel function."""
    # Warmup
    for _ in range(warmup):
        fn(*args)
    torch.cuda.synchronize()

    start = time.perf_counter()
    for _ in range(iterations):
        fn(*args)
    torch.cuda.synchronize()
    end = time.perf_counter()

    avg_ms = (end - start) / iterations * 1000
    return avg_ms

# Example usage
input_tensor = torch.randn(32, 2048, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(2048, dtype=torch.bfloat16, device="cuda")

pytorch_time = benchmark_kernel(
    torch.nn.functional.rms_norm, input_tensor, (2048,), weight, 1e-6
)
custom_time = benchmark_kernel(
    cuda_rmsnorm, input_tensor, weight, 1e-6
)

print(f"PyTorch: {pytorch_time:.3f} ms")
print(f"Custom:  {custom_time:.3f} ms")
print(f"Speedup: {pytorch_time / custom_time:.2f}x")

nsys Profiling

nsys profile --stats=true \
    --trace=cuda,nvtx \
    -o h100_diffusers_profile \
    python run_pipeline.py

Common Performance Issues

Symptom	Likely Cause	Fix
No speedup over PyTorch	Kernel launch overhead dominates	Fuse operations, use larger batch sizes
Slower than PyTorch	Bank conflicts in shared memory	Pad shared memory arrays
Inconsistent results	Race condition	Check synchronization barriers
NaN outputs	Overflow in BF16	Add epsilon before division, check ranges
Wrong results on some inputs	Edge case in vectorized loads	Handle non-aligned tail elements

Summary

• Custom CUDA kernels provide 6% standalone and 34% combined with torch.compile end-to-end speedup
• RMSNorm micro-benchmarks show 2.67x speedup over PyTorch
• Always handle the four critical pitfalls: None weights, diffusers vs torch RMSNorm, GELU vs GEGLU per model, and injection ordering with CPU offloading
• Register kernels as custom ops for torch.compile compatibility
• Target H100's 132 SMs and 3.35 TB/s bandwidth with vectorized memory access patterns