PyTorch Custom Operators: Schema, Fake Kernels, and Autograd

Core Concept

A kernel defines how the GPU computes. An operator also defines schema, mutation and aliasing, dispatch, fake metadata, backward behavior, and compiler integration.

Why It Affects Performance

Incorrect contracts allow unsafe reordering or reuse. Missing fake implementations block compile/export. Registration does not make a kernel faster and can become a fusion boundary.

Example or Pseudocode

@torch.library.custom_op('gpu_lessons::saxpy', mutates_args=(), device_types='cuda')
def saxpy_op(x: torch.Tensor, y: torch.Tensor, a: float) -> torch.Tensor: ...
@saxpy_op.register_fake
def _(x, y, a): return torch.empty_like(x)
def backward(ctx, g): return g * ctx.a, g, None

Profiling and Performance Judgment

Use torch.library.opcheck for registration contracts, plus independent forward, gradient, boundary, dtype, and layout tests. Profile eager, compile, direct Triton, and registered versions.

Five-Minute Exercise

For z=ax+y, what must the custom backward save?