Skip to content

Extensions in Practice & Pitfalls

Overview

Writing the extension is the easy part; integrating it cleanly is where projects die. This file collects the practical lessons: build hygiene, gradient correctness, determinism, debugging, and when to say "don't write a custom op."

Practice Checklist

  1. Keep the C++ minimal— register only the kernel; do shape logic in Python.
  2. Guard everything— TORCH_CHECK on shapes, dtypes, contiguity, devices.
  3. Test parity— reference check vs a pure-Python implementation on random inputs (seeded).
  4. Wrap in autograd.Function with a torch.autograd.gradcheck pass.
  5. Profile before and after— a custom op should beat the Python path on your real workload.
  6. Build reproducibly— pin torch version; rebuild after upgrades.
import torch

# the full test sandwich
from torch.autograd import gradcheck

class MyFn(torch.autograd.Function):
 @staticmethod
 def forward(ctx, x):
 return myops.fn(x) # your compiled op
 @staticmethod
 def backward(ctx, g):
 return g * 2 # hand-derived gradient

x = torch.randn(3, 3, dtype=torch.double, requires_grad=True)
assert gradcheck(MyFn.apply, (x,), eps=1e-6, atol=1e-4)
print("gradcheck passed ")

-

Gradient Correctness— The Two Failure Modes

  1. Wrong math in backward → silent divergence. Always gradcheck + manual finite-difference spot check.
  2. Wrong arity / None placement → runtime errors. Return one grad per forward arg, None for constants.
# spot finite-difference
x = torch.tensor([2.0], requires_grad=True)
h = 1e-5
fd = (fn(x + h) - fn(x - h)) / (2 * h)
an = torch.autograd.grad(fn(x), x)[0]
print("fd vs autograd:", fd.item(), an.item(), "close:", torch.allclose(fd, an, atol=1e-4))

Determinism

Concern Fix
Same inputs → same outputs avoid atomics; set deterministic flags
Same seed across runs torch.manual_seed + torch.use_deterministic_algorithms(True)
Reductions ordering vectorize or use fixed block reduction
Random in kernel curand seeded per launch— still deterministic if seeded 
torch.use_deterministic_algorithms(True)
# will raise if any non-deterministic op is used

Debugging Cookbook

Symptom Likely cause Tool
Crash on certain shapes OOB index compute-sanitizer --tool memcheck
Wrong values only on big tensors race / uninitialized ncu, nsight, unit tests at multiple sizes
Works in python, fails after build stale .so rebuild clean; -D_GLIBCXX_USE_CXX11_ABI=1
Non-determinism atomics/races compute-sanitizer --tool racecheck
Slow despite "custom" naive kernel profile; consider Triton

-

Build & Distribution Pitfalls

  1. Torch version drift— rebuild after every torch upgrade (ABI changes).
  2. JIT cache growth— load() caches per (name, flags); clean ~/.cache/torch_extensions when weird.
  3. Windows/Linux flags differ— MSVC vs GCC; keep flags minimal.
  4. Multi-extension module— one .so per logical op family; name uniquely.
  5. Linking CUDA-only code without GPU present— build with with_cuda=True but guard runtime.

-

When NOT to Write a Custom Op

- Elementwise chains → torch.compile / Triton
- Attention → F.scaled_dot_product_attention (fused, battle-tested)
- Reductions → torch built-ins + Inductor
- Rare one-off math → torch + autograd is fine
- Learning exercise → write it, then replace with the built-in

The strongest signal: after profiling, if the op isn't the bottleneck, don't write it.

Production Integration

  • Expose via a thin Python API with type hints & docstrings.
  • Add to your CI: compile + gradcheck + parity test on CPU and GPU.
  • Pin versions; record build flags in README.
  • Fallback path: try: import myops except ImportError: use pure-python fallback— keeps training servers alive on bare-metal.
try:
 import myops
 USE_EXT = True
except ImportError:
 USE_EXT = False

def fast_add(a, b):
 return myops.add(a, b) if USE_EXT else a + b

-

Key Takeaways

  • Correctness gate: gradcheck + finite-difference + parity tests on every change.
  • Determinism matters for reproducibility— test and pin it.
  • Debug OOB/races with compute-sanitizer; profile with ncu before claiming speed.
  • Rebuild on torch upgrades; keep JIT caches clean; document flags.
  • Most "custom op" urges are served better by torch.compile/Triton/built-ins— prove the need first.

-