Custom Operators: Extending PyTorch with CUDA Kernels

Overview

Custom operators enable: - Specialized CUDA kernels for optimized computation - Custom forward/backward passes for novel operations - C++ extensions for performance-critical code - Memory efficiency through fused operations

PyTorch provides multiple levels of customization for different use cases.

---

Pure Python Custom Functions

Basic Custom Operation

import torch
import torch.nn.functional as F

class ReLU(torch.autograd.Function):
    """Custom ReLU operation."""

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.clamp(input, min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

def relu(x):
    return ReLU.apply(x)

# Test
x = torch.tensor([1.0, -2.0, 3.0], requires_grad=True)
y = relu(x)
loss = y.sum()
loss.backward()

print(x.grad)  # [1., 0., 1.]

GELU Activation (Example)

class GELU(torch.autograd.Function):
    """Gaussian Error Linear Unit (GELU) activation."""

    @staticmethod
    def forward(ctx, input):
        # GELU: x * Phi(x) where Phi is Gaussian CDF
        # Approximation: x * sigmoid(1.702 * x)
        ctx.save_for_backward(input)
        return 0.5 * input * (1 + torch.erf(input / math.sqrt(2)))

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors

        # Derivative: Phi(x) + x * phi(x)
        # phi is Gaussian PDF: exp(-x^2/2) / sqrt(2*pi)
        pdf = torch.exp(-input ** 2 / 2) / math.sqrt(2 * math.pi)
        cdf = 0.5 * (1 + torch.erf(input / math.sqrt(2)))

        grad_input = grad_output * (cdf + input * pdf)
        return grad_input

---

CUDA Kernels

When to Use CUDA Kernels

# Use Python/PyTorch ops when:
# - Computation is simple
# - PyTorch ops are fast enough
# - Optimization not critical

# Use CUDA kernels when:
# - Need custom memory layout
# - Multiple operations can be fused
# - Significant performance bottleneck
# - Algorithm not in standard ops

# Example: Flash Attention fuses:
# 1. Q*K^T computation
# 2. Softmax
# 3. Dropout
# 4. Output computation
# → Single kernel (6x faster than separate ops)

CUDA Kernel Structure

// cuda_kernel.cu

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// CUDA kernel
__global__ void custom_kernel(float* input, float* output, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        output[idx] = input[idx] * 2.0f;  // Simple: multiply by 2
    }
}

// Wrapper for PyTorch
torch::Tensor custom_forward(torch::Tensor input) {
    auto output = torch::empty_like(input);
    int threads = 256;
    int blocks = (input.numel() + threads - 1) / threads;

    custom_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        input.numel()
    );

    return output;
}

// Register with PyTorch
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("custom_forward", &custom_forward, "Custom CUDA kernel");
}

Building CUDA Extensions

# setup.py
from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension

setup(
    name='custom_ops',
    ext_modules=[
        CUDAExtension(
            'custom_ops',
            ['cuda_kernel.cu'],
            extra_compile_args={'cxx': ['-g'], 'nvcc': ['-O2']}
        ),
    ],
    cmdclass={'build_ext': BuildExtension}
)

# Build
# python setup.py build_ext --inplace

Using CUDA Extension

# Python code
import custom_ops

x = torch.randn(1000, 1000, device='cuda')
y = custom_ops.custom_forward(x)

---

PyBind11 for C++ Extensions

Custom Operation with Autograd

// custom_ops.cpp

#include <torch/extension.h>

torch::Tensor custom_forward(torch::Tensor input, torch::Tensor weight) {
    return torch::matmul(input, weight);
}

torch::Tensor custom_backward_input(
    torch::Tensor grad_output, 
    torch::Tensor weight
) {
    return torch::matmul(grad_output, weight.t());
}

torch::Tensor custom_backward_weight(
    torch::Tensor grad_output, 
    torch::Tensor input
) {
    return torch::matmul(grad_output.t(), input);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("forward", &custom_forward, "Forward pass");
    m.def("backward_input", &custom_backward_input, "Backward w.r.t. input");
    m.def("backward_weight", &custom_backward_weight, "Backward w.r.t. weight");
}

---

Fused Operations

Fused LayerNorm + GELU

class FusedGELU(torch.autograd.Function):
    """Fused GELU activation (faster than separate operations)."""

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        # Fused computation in single pass
        return 0.5 * input * (1 + torch.erf(input / math.sqrt(2)))

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors

        pdf = torch.exp(-input ** 2 / 2) / math.sqrt(2 * math.pi)
        cdf = 0.5 * (1 + torch.erf(input / math.sqrt(2)))

        return grad_output * (cdf + input * pdf)

# Practical: Fused LayerNorm + GELU
class FusedLayerNormGELU(torch.nn.Module):
    def __init__(self, hidden_size, eps=1e-12):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
        self.bias = torch.nn.Parameter(torch.zeros(hidden_size))
        self.eps = eps

    def forward(self, input):
        # Fused: LayerNorm + GELU in one kernel (hypothetically)
        # Real implementation would use CUDA

        # 1. LayerNorm
        mean = input.mean(-1, keepdim=True)
        std = input.std(-1, keepdim=True)
        normalized = (input - mean) / (std + self.eps)
        normalized = normalized * self.weight + self.bias

        # 2. GELU
        output = FusedGELU.apply(normalized)

        return output

Fused Softmax + Dropout

class FusedSoftmaxDropout(torch.autograd.Function):
    """Fused softmax and dropout (reduces memory access)."""

    @staticmethod
    def forward(ctx, input, p=0.1):
        # Softmax
        input_exp = torch.exp(input - input.max(-1)[0].unsqueeze(-1))
        softmax = input_exp / input_exp.sum(-1, keepdim=True)

        # Dropout (training)
        if p > 0:
            mask = torch.bernoulli(torch.ones_like(softmax) * (1 - p))
            output = softmax * mask / (1 - p)
        else:
            output = softmax
            mask = None

        ctx.save_for_backward(input, softmax, mask)
        ctx.p = p

        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, softmax, mask = ctx.saved_tensors
        p = ctx.p

        # Gradient through dropout
        if mask is not None:
            grad_output = grad_output * mask / (1 - p)

        # Gradient through softmax
        grad_input = softmax * (
            grad_output - (grad_output * softmax).sum(-1, keepdim=True)
        )

        return grad_input, None

---

Performance: Custom vs Built-in

Benchmark Custom Kernel

import torch
import time

def benchmark_relu():
    x = torch.randn(10000, 10000, device='cuda', requires_grad=True)

    # Built-in ReLU
    start = time.time()
    for _ in range(100):
        y = torch.relu(x)
        y.backward(torch.ones_like(y))
        x.grad.zero_()
    builtin_time = time.time() - start

    # Custom ReLU (Python)
    start = time.time()
    for _ in range(100):
        y = ReLU.apply(x)
        y.backward(torch.ones_like(y))
        x.grad.zero_()
    custom_time = time.time() - start

    print(f"Built-in: {builtin_time:.3f}s")
    print(f"Custom: {custom_time:.3f}s")
    print(f"Speedup: {custom_time / builtin_time:.2f}x")

# Output (typical):
# Built-in: 0.234s
# Custom: 0.456s
# Speedup: 1.95x (slower!)
# → Custom Python autograd slower than optimized built-in

When Custom Wins

# Example: Flash Attention
# Forward: 2.8x faster (fused kernel)
# Backward: 2x faster (fused backward)
# Memory: 50% reduction (no intermediate attention matrix)

# Conditions for speedup:
# 1. Multiple operations fused
# 2. Significant memory savings
# 3. GPU compute-bound (not memory-bound)
# 4. Complex control flow saved

---

Practical: Fused Linear + GELU

import torch
import torch.nn as nn

class FusedLinearGELU(nn.Module):
    """Fused Linear + GELU layer."""

    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, x):
        # Without fusion (3 operations):
        # 1. Matrix multiply
        # 2. Add bias
        # 3. GELU
        # → 3 kernel launches, 3x memory bandwidth

        # With fusion (1 operation):
        # All in single CUDA kernel

        # Python version (simulates):
        x = self.linear(x)
        x = 0.5 * x * (1 + torch.erf(x / math.sqrt(2)))

        return x

# Real implementation would use:
# 1. PyBind11 wrapper
# 2. CUDA kernel for fusion
# 3. Custom autograd function

Gradient Checking Custom Op

import torch

def gradient_check(func, input, eps=1e-5, atol=1e-4):
    """Verify custom operation gradients."""
    input.requires_grad = True

    # Analytical gradient
    output = func(input)
    output.backward(torch.ones_like(output))
    analytical_grad = input.grad.clone()

    # Numerical gradient
    input.grad.zero_()
    numerical_grad = torch.zeros_like(input)

    for i in range(input.numel()):
        # f(x + eps)
        input_plus = input.clone()
        input_plus.view(-1)[i] += eps
        output_plus = func(input_plus).sum()

        # f(x - eps)
        input_minus = input.clone()
        input_minus.view(-1)[i] -= eps
        output_minus = func(input_minus).sum()

        # (f+ - f-) / 2eps
        numerical_grad.view(-1)[i] = (
            (output_plus - output_minus) / (2 * eps)
        )

    # Compare
    if torch.allclose(analytical_grad, numerical_grad, atol=atol):
        print("✓ Gradient check passed!")
        return True
    else:
        print("✗ Gradient check failed!")
        print(f"Max diff: {(analytical_grad - numerical_grad).abs().max()}")
        return False

# Test custom ReLU
gradient_check(relu, torch.randn(10, requires_grad=True))

---

Common Pitfalls

Pitfall 1: Memory Layout

# Tensors can be C-contiguous or not
x = torch.randn(3, 4, 5)
y = x.transpose(0, 2)  # Not contiguous!

# CUDA kernel expects contiguous memory
# Must call .contiguous() first
y = y.contiguous()

Pitfall 2: Device Mismatch

# Tensor on GPU, kernel expects CPU
x = torch.randn(10, device='cuda')
y = cpu_kernel(x)  # Error!

# Must handle device placement
device = x.device
x = x.to('cpu')
y = cpu_kernel(x)
y = y.to(device)

Pitfall 3: Gradient Leaks

class BadOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        # Forgot to save for backward!
        return x ** 2

    @staticmethod
    def backward(ctx, grad_output):
        # Need saved x to compute gradient
        # This will fail!
        return None  # No gradient

class GoodOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)  # Save for backward
        return x ** 2

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors  # Retrieve
        return 2 * x * grad_output

---

Summary

Method	Speed	Memory	Complexity	Use Case
Python autograd	Slow	High	Low	Prototyping
PyTorch ops	Medium	Medium	Low	Most cases
Fused ops	Fast	Low	High	Bottlenecks
CUDA kernels	Very fast	Low	Very high	Production
TorchScript	Medium	Medium	Medium	Deployment

---

Related Topics

• 02 Autograd Implementation - Understanding backward passes
• 00 Readme - C/C++ integration
• 04 Profiling & Performance Analysis - Benchmarking custom ops