PyTorch custom operator requirements
Define schema, fake-tensor, autograd, autocast, dispatch, and fallback behavior for a custom operation
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- PyTorch tensors and autograd
- C++/CUDA build concepts
- Operator semantics
A developer imports a compiled CUDA extension and calls its fused operation from PyTorch. The call returns correct values in eager mode for one contiguous tensor. Running torch.compile causes a graph break, calling backward produces no gradient, and a transposed input returns incorrect values because the CUDA code assumed contiguous storage. PyTorch operator checks and compiler logs make each failure visible.
Registering a custom operator gives PyTorch a stable name and a precise contract for arguments, outputs, mutation, aliasing, shapes, datatypes, and differentiation. A fake implementation describes output metadata for tracing without reading real tensor data. An autograd registration describes the backward calculation. Device implementations perform actual work. The separate pieces let PyTorch reason about the operation before a GPU kernel executes.
A custom operator contract specifies output shapes and datatypes, device behavior, aliasing, mutation, errors, gradients, autocast, and compatibility with graph capture. The kernel implements part of this contract. An incomplete contract can work in eager forward execution and fail during compilation, export, fake-tensor execution, distributed wrapping, or backward.
Framework integration can remove a kernel-level improvement by creating graph breaks or materializing conversions around the call. Define the operator schema, dispatch region, and compiler behavior while you design the kernel.
Treat the reference implementation and operator schema as the stable interface. CUDA or Triton kernels are replaceable implementations. Separating interface from implementation lets tests detect semantic changes and lets the framework use a correct fallback when a specialization is not eligible.
Define the operator contract
The operator schema is its stable framework signature. The schema names the operator, its argument and result types, defaults, and mutation or aliasing behavior. State input ranks, shapes, devices, datatypes, strides, and legal aliasing in the surrounding contract as well. PyTorch transformations rely on this information for functionalization, which rewrites mutation into functional equivalents, and for autograd, fake tensors, and compilation. An incorrect mutation declaration can lead to invalid reordering or reused values.
State numerical details that ordinary function names hide. For normalization, include epsilon placement, accumulation datatype, output datatype, treatment of non-finite inputs, and whether a residual output is returned. The high-level reference should implement exactly this contract and remain available for unsupported devices or shapes.
Provide fake-tensor behavior
A fake tensor contains metadata such as shape, stride, datatype, and device but has no ordinary data storage. Compiler and export systems use fake execution to determine output metadata without running a device kernel. The fake implementation must create outputs with the correct metadata and must not inspect data values.
If output shape depends on input values, a simple fake implementation cannot read those values to determine the shape. The operator may need a symbolic size mechanism, a constrained contract, or a different interface that returns capacity plus a valid length. Test fake execution with dynamic shapes because an eager-only test does not exercise this path.
Register gradients and mixed-precision behavior
A differentiable operator needs a backward formula. The formula can call a custom backward kernel or express the gradient with existing PyTorch operations. Decide which forward values to save and which to recompute. Saving large tensors increases activation memory; recomputation increases arithmetic. The choice should match the operator's training objective.
Use gradcheck with a datatype and tolerances that can expose mathematical errors. Then add lower-precision tests for the actual deployment types. Autocast is the framework policy that runs eligible operations in selected lower-precision formats while retaining wider formats where required. The custom operator must specify which inputs are converted, which precision performs reductions, and which datatype the output uses. Do not let incidental wrapper conversions define the contract.
Make dispatch conditions explicit
The optimized kernel may require contiguous rows, aligned pointers, a bounded hidden dimension, a supported architecture, or a specific datatype. Test these requirements in one eligibility function. Inputs that fail should reach the reference or another correct kernel. Silent reinterpretation of unsupported strides is a correctness defect.
Dispatch has measurable overhead for very small operations. Layout conversions inserted before the call also count against the optimization. Benchmark from the original input layout through the final required output, and exercise input cases just inside and just outside every specialized region.
Choose the compiler boundary
An opaque custom operation lets compilation retain a call without tracing into its implementation, provided the framework knows its metadata and effects. The opaque operation is useful for native code that Dynamo cannot inspect. The operator interface can prevent fusion or layout propagation across the call.
A decomposition exposes the operation as existing framework operations, which enables compiler transformations but may bypass the specialized kernel. PyTorch also provides structured integration for user-defined Triton kernels in compilation. Inspect the captured and lowered graphs to confirm which path executes. A successful compile call does not prove that the optimized kernel remained in the graph or that surrounding conversions were removed.
Worked example: Integrating a fused RMSNorm operator
A CUDA or Triton kernel fuses residual addition and RMSNorm. The fused operator must work in eager training and compiled inference.
- Write a trusted PyTorch reference that defines outputs, residual behavior, epsilon, datatype, and supported strides. Use it for correctness comparisons and as a fallback implementation.
- Declare mutation and aliasing. For a functional RMSNorm result, returned tensors must be fresh values unless the schema explicitly promises an alias.
- Register the operator schema and fake implementation so symbolic shapes produce correct metadata without allocating device data.
- Provide backward or a tested decomposition for training. Decide which tensors to save and validate gradients across dtypes and odd hidden widths.
- Compile an enclosing model and inspect graph breaks, inserted conversions, and fusion boundaries. Benchmark eager and compiled end to end across the actual shape distribution.
- Run torch.library.opcheck for registration behavior and gradcheck for the mathematical gradient. The two tests answer different questions and neither replaces numerical forward tests.
Conclusion: The deliverable is an operator that remains fast and correct across framework modes, not a kernel reachable only from one benchmark script.
Common errors
- Treating contiguous inputs as an unstated assumption. Either enforce it visibly, handle strides, or dispatch to a fallback.
- Registering only forward. Training, fake execution, export, and autocast can fail far from the original call.
- Assuming an opaque node is compiler-friendly. The opaque node may preserve capture while blocking layout or fusion opportunities around the call.
- Making the fake implementation read tensor data. Fake tensors do not provide ordinary storage and compiler shape reasoning cannot depend on runtime values this way.
Reference summary
A production PyTorch operator needs a schema: the stable operator name, argument and return types, defaults, and mutation or aliasing contract. The operator also needs device implementations, shape, stride, and datatype validation, fake or meta behavior for tracing, autograd behavior when differentiable, explicit mixed-precision behavior, and tests that cover compilation. Directly calling an extension can appear to work while creating graph breaks, wrong gradients, or unsafe behavior on non-contiguous inputs.
- Define semantics independently of one kernel implementation.
- Register CPU/reference and CUDA implementations where appropriate.
- Provide fake-tensor shape propagation without reading real data.
- Register autograd or a custom backward and test with gradcheck at suitable precision.
- Run framework operator checks plus adversarial shape/stride/dtype tests.
- Benchmark dispatch and launch overhead for small operators.
Knowledge check
What information must a fake/meta implementation provide, and why must it avoid reading tensor data?
Show the answer guide
- A fake or meta implementation must compute output metadata from input metadata: output shape, dtype, device class, layout or strides when relevant, and any aliasing or mutation relationship required by the operator contract.
- The implementation must also reproduce shape validation and metadata-dependent errors so graph capture can reason about valid programs before real execution.
- Fake or meta tensors may have no real storage and intentionally do not contain usable element values. Reading tensor data would fail during abstract execution or make the graph depend on values that are unavailable until runtime.
- If output shape genuinely depends on tensor values, the operator needs a declared dynamic-output mechanism or a redesigned interface; the meta function must not guess by executing the real kernel.
Define and test a custom operator's abstract contract
Register a custom PyTorch row-normalization operator for inputs shaped [B,N] that returns a normalized tensor and one FP32 scale per row. Provide CPU or reference behavior, the accelerated implementation, and a fake or meta implementation derived only from shapes, dtypes, devices, and strides. Exercise contiguous, transposed, empty-row, FP16, and BF16 inputs through eager execution and torch.compile.
Evidence to keep: Submit registration code, fake or meta code, eager correctness and gradient tests, compile or export tests using fake tensors, and a captured graph or compiler log that shows the output metadata. Include a failing test that proves an attempted data read in abstract execution is invalid, without retaining the invalid implementation.
A shared upstream operator requires tests, benchmarks, documentation, maintainable code, and support for the declared input range.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.