Custom Operators

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
# 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
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
# Custom
# Speedup
# → Custom Python autograd slower than optimized built-in

When Custom Wins

# Example
# Forward
# Backward
# Memory

# 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](/05-py3/09-bytecode-and-execution/(04-profiling-performance-analysis/) - Benchmarking custom ops