TopolearnCustom kernels (Triton, CUDA, C++) only integrate cleanly with compile, export and autograd through the operator registry.
Every torch.* call ends up in the dispatcher, which routes an operator such as aten::add to an implementation based on dispatch keys computed from the inputs: Autograd, then backend keys like CPU, CUDA or Meta, plus features like Autocast or Python subclasses. That is how one operator gets autograd, mixed precision and device-specific kernels without the caller knowing.
Python can hook in at two levels:
__torch_function__on a Tensor subclass intercepts calls at the Python API level.__torch_dispatch__intercepts after autograd, at the ATen operator level; it is how fake tensors, functionalisation and DTensor are built.
To add your own operator, use `torch.library.custom_op`. It registers the op with a schema inferred from type hints, and you add:
register_fake: a shape-and-dtype-only implementation, sotorch.compileandtorch.exportcan trace through the op without running it.register_autograd(backward, setup_context=...): the derivative.torch.library.opcheck(op, args): tests that the registrations are consistent.
The implementation can call anything: a Triton kernel, a C++ extension, even NumPy. Custom ops must not return views of their inputs, and they appear as single opaque nodes in compiled graphs. For C++ and CUDA kernels, torch.utils.cpp_extension builds them and TORCH_LIBRARY registers them.
Example
import torch
from torch import Tensor
@torch.library.custom_op("mylib::scaled_tanh", mutates_args=())
def scaled_tanh(x: Tensor, scale: float) -> Tensor:
return torch.tanh(x) * scale # could call Triton, C++ or NumPy here
@scaled_tanh.register_fake
def _(x, scale):
return torch.empty_like(x) # shapes/dtypes only, used while tracing
def setup_context(ctx, inputs, output):
_, scale = inputs
ctx.save_for_backward(output)
ctx.scale = scale
def backward(ctx, grad):
(y,) = ctx.saved_tensors
return grad * ctx.scale * (1 - (y / ctx.scale) ** 2), None
scaled_tanh.register_autograd(backward, setup_context=setup_context)
x = torch.randn(4, requires_grad=True)
torch.ops.mylib.scaled_tanh(x, 2.0).sum().backward()
print(x.grad.shape) # torch.Size([4])Remember
- The dispatcher routes ops by keys: Autograd, Autocast, CPU, CUDA, Meta...
torch.library.custom_op+register_fake+register_autograd- Test registrations with
torch.library.opcheck
Watch out for
- Skipping
register_fake, which makes torch.compile graph-break or fail on the op.