Validating Pytorch Custom Ops

Validate PyTorch Implementations and Custom Operators

Use this workflow whenever adding a backend under attn_gym/linear/<variant>/impl/ or attn_gym/sparse/<variant>/impl/.

New operator checklist

1. Pick the integration below; for an opaque kernel, register a fixed-arity define/impl pair plus an autograd wrapper.
2. Write each fake from the launcher's real allocations (see the fake rules below).
3. opcheck with compact inputs and a dense permuted layout of each tensor input (p.t().contiguous().t(), or token-innermost x.movedim(1, -1).contiguous().movedim(-1, 1)). Padded or sliced inputs do not substitute: empty_like of a non-dense tensor is already contiguous, so a stride-copying fake still passes.
4. Validate numerics and fullgraph=True compile on every supported architecture; Hopper and Blackwell often dispatch backends that allocate differently.

Choose the least opaque integration

1. Eager/reference implementation: keep ordinary PyTorch code traceable. Do not create a custom operator merely because the public function must support torch.compile.
2. Triton implementation: direct Triton kernels can work under torch.compile. Prefer torch.library.triton_op with torch.library.wrap_triton when a stable operator boundary is useful and compiler subsystems should remain able to inspect the implementation.
3. CuTeDSL or external implementation: use a private registered operator (torch.library.define/impl, below) when the launcher cannot be traced correctly.

The documented public function remains an ordinary Python function that owns semantic validation, mode selection, and backend dispatch. Compiling only a private launcher does not prove that the public operation is compile-friendly.

Operator boundary CPU overhead

Every registered-operator boundary costs CPU time per call. Microbenchmark (torch 2.14 nightly, B200, identical trivial launcher, CPU wall-clock per call, medians over interleaved rounds):

Nesting one registered op inside another pays the boundary again (nested custom_op 7.1 us, nested define/impl 3.7 us). Wrapping the real kda_l2norm_fwd launcher in custom_op added ~11 us per forward call over invoking the launcher directly.

Consequences for launch-bound paths (small kernels, high call rate):

• Registration exists for compile and fake-tensor consumers; eager execution gains nothing from it. The eager floor is a plain autograd.Function over the launcher with no torch.library registration (real l2norm op: ~17 us cheaper per gradient-tracking forward and ~58 us per fwd+bwd than the custom_op boundary), at the cost of fullgraph=True support.
• Use the define/impl pattern below for repo operators; torch.library.custom_op costs ~3 us more per call for schema inference and mutation checking and is only preferable for prototypes.
• Prefer autograd.Function over register_autograd for hot training ops; it was the cheapest measured autograd boundary. Registering a Python kernel at the Autograd dispatch key (torch.library.impl(op, "Autograd")) makes the raw op differentiable but runs Python on every call, including inference; use it only when direct torch.ops calls must be differentiable.
• Give each hot path at most one registered-operator boundary. Nested custom ops double the dispatch tax; share the Python launcher between entrypoint ops instead of calling one op from another.
• One contract, one operator pair. When several backends (Triton portable, CuTeDSL fast) or variants (a transform kind, a mode) implement the same public semantics, register a single fwd/bwd pair, pass the variant as a schema argument (str kind, bool, float), and select the backend inside the CUDA implementation on real tensors. Do not create one op pair per backend or thread a use_<backend> flag through autograd.Function.apply: selection that runs under Dynamo tracing specializes into the graph, trips the functools.lru_cache warning for cached device-property helpers, and multiplies fakes, handles, and Functions. Make every backend return the same output contract (e.g. already-reduced parameter gradients) so one fake describes all of them. attn_gym::_gate_transform_{fwd,bwd} is the reference shape.
• A define/impl op has no autograd kernel: backprop through a direct call warns and produces no gradient, so route all differentiable use through the autograd.Function wrapper.
• torch.library.opcheck accepts torch.ops.attn_gym.op.default, so the validation workflow below is unchanged.

Public result contract

Private backends should return tensors or a fixed tensor tuple. The public API may wrap that result in a documented NamedTuple so users can use attributes or unpack it:

from typing import NamedTuple

import torch


class OperationOutput(NamedTuple):
    output: torch.Tensor
    final_state: torch.Tensor | None = None


def operation(...) -> OperationOutput:
    output, final_state = backend_forward(...)
    return OperationOutput(output, final_state)

A registered operator's output structure and arity must agree with its schema. Do not switch between returning a tensor and a tuple based on an argument. Represent optional outputs explicitly and match them in the fake implementation.

Opaque operator pattern (define/impl + autograd.Function)

The repo standard for opaque kernel boundaries. Write the schema string yourself; nothing is inferred from annotations:

import torch
from torch import Tensor

torch.library.define(
    "attn_gym::example_fwd",
    "(Tensor query, Tensor key, Tensor value, Tensor? initial_state) -> (Tensor, Tensor)",
)
torch.library.define("attn_gym::example_bwd", "(Tensor query, Tensor grad_output) -> Tensor")


def _example_fwd_cuda(
    query: Tensor, key: Tensor, value: Tensor, initial_state: Tensor | None
) -> tuple[Tensor, Tensor]:
    return launch_backend(query, key, value, initial_state)


# Register with an explicit call, not as a decorator: the decorator form returns None,
# and keeping the launcher callable lets benchmarks and tests bypass the boundary.
torch.library.impl("attn_gym::example_fwd", "CUDA", _example_fwd_cuda)


@torch.library.register_fake("attn_gym::example_fwd")
def _example_fwd_fake(
    query: Tensor, key: Tensor, value: Tensor, initial_state: Tensor | None
) -> tuple[Tensor, Tensor]:
    return value.new_empty(value.shape), query.new_empty(query.shape[:2])


_example_fwd = torch.ops.attn_gym.example_fwd.default
_example_bwd = torch.ops.attn_gym.example_bwd.default  # impl/fake registered the same way


class _Example(torch.autograd.Function):
    @staticmethod
    def forward(ctx, query: Tensor, key: Tensor, value: Tensor) -> Tensor:
        output, state = _example_fwd(query, key, value, None)
        ctx.save_for_backward(query)
        return output

    @staticmethod
    @torch.autograd.function.once_differentiable
    def backward(ctx, grad_output: Tensor):
        (query,) = ctx.saved_tensors
        return _example_bwd(query, grad_output), None, None

Schema rules:

• Optional tensors are Tensor?; scalars are int, float, bool; defaults may be embedded (Tensor? cu_seqlens=None).
• Encode mutation explicitly with alias annotations: (Tensor(a!) state, Tensor value) -> ().
• The schema language is a superset of what custom_op annotation inference accepts, adding e.g. ScalarType, Layout, MemoryFormat, Generator, SymInt, and int[2]. Neither flow accepts arbitrary Python objects, dataclasses, or callables; flatten configs to scalars or specialize per-config outside the op.
• Optional outputs parse as Tensor? and, as of torch 2.15 nightly (verified 2026-08-16: plain fullgraph compile, autograd.Function wrapping + backward, and opcheck all pass for both the tensor and None branches), no longer hard-fail the compile stack; on older releases the None return handling was unreliable. Keep the fixed-arity convention anyway: for "N or N+1 returns depending on a flag", define two fixed-arity schemas sharing one launcher (see kda_chunk_fwd / kda_chunk_fwd_with_state) and let the autograd.Function branch on the flag; its output arity is free to vary. The bool flag specializes into separate graphs regardless, so a merged optional-output schema saves nothing under compile, while a Tensor? return forces None-narrowing on every caller and extra schema boxing costs real time on launch-bound eager paths (collapsing schema pairs measured 3.22% slower on the dense forward probe, PR #314).
• Optional inputs are fully compile-safe: the merged kda_chunk_bwd schema takes Tensor? cu_seqlens, Tensor? chunk_offsets for both dense and ragged and passes the strict fullgraph matrix. The forward keeps separate dense/ragged ops as a measured eager-dispatch decision, not a compile requirement.
• Tensors needed only by the backward (autograd tapes, packing metadata) should be op outputs saved by the autograd.Function via ctx.save_for_backward, not part of the wrapper's user-facing return. Unlike register_autograd, the Function is not limited to saving user-visible outputs.
• Bind torch.ops.attn_gym.<op>.default to a module-level name and call that; resolving the torch.ops attribute chain per call adds overhead.

The fake implementation describes output metadata without running the kernel. It must:

• return the same output structure as the real implementation;
• preserve shape, dtype, device, layout, and relevant strides. Allocate outputs contiguously in both the launcher and the fake (x.new_empty(shape) / torch.empty(...)), not with torch.empty_like(x): it copies a dense permuted input's strides, so the fake can disagree with the kernel (Inductor assert_size_stride failure) and a kernel indexing outputs with contiguous offsets writes the wrong elements;
• avoid reading tensor values, storage, data pointers, or calling .item();
• use symbolic shape arithmetic rather than data-dependent Python branches.

For genuinely data-dependent output dimensions, use torch.library.get_ctx() and new_dynamic_size() rather than inspecting input data.

Registration must exist before graph capture. Keep registration deterministic and lightweight; do not rely on registration side effects occurring inside a compiled region. Continue importing optional kernel dependencies lazily.

CuTeDSL layout and alignment contracts

Expose the weakest layout contract the kernel actually needs. Do not require full contiguity when only one mode must be contiguous, and do not treat alignment as a runtime repair mechanism: assumed_align is a promise to codegen that may enable wide loads.

Use the shared helpers in attn_gym._backends.cute:

• tensor_supports_contiguous_dim(tensor, dim=-1, alignment_bytes=...) checks that one mode has stride 1 and that every slice origin satisfies the promised byte alignment;
• make_fake_strided_tensor(dtype, shape, contiguous_dim=-1, ...) creates the matching TVM-FFI fake signature with dynamic other strides and honest minimum alignment;
• tensor_supports_tma(tensor) is the 16-byte CUDA/TMA specialization of the same predicate;
• requires_int64_abi(...) still selects the independent wide-stride/address specialization.

Prefer two measured signatures when vectorization needs stronger assumptions:

1. a compact/aligned fast signature advertising the vector width;
2. a general signature with the required contiguous mode, dynamic outer strides, and only element-size alignment.

Normalize only unsupported layouts. A noncontiguous required mode may need a compact copy, but an outer-strided or misaligned tensor belongs on the general signature if the generated loads are safe. Note that .contiguous() returns an already-contiguous tensor unchanged and therefore does not fix a misaligned storage offset.

For every new CuTe tensor ABI, test:

• compact aligned input selects the fast path;
• arbitrary aligned outer strides with the required contiguous mode;
• a noncontiguous required mode and its documented fallback or rejection;
• contiguous storage with a deliberately misaligned offset;
• partial logical tiles (odd token counts and non-divisible head/channel groups);
• every supported dtype, fake/opcheck, dynamic fullgraph reuse, and forced-int64 equivalence.

Express views through native slicing and local_tile/partition algebra. Reconstruct a tensor from .iterator only for a deliberate storage reinterpretation such as a TMA mode permutation, and explain that contract at the call site.

Mutation and aliasing

Declare every mutated argument in the schema string:

torch.library.define("attn_gym::update_state", "(Tensor(a!) state, Tensor value) -> ()")

Do not declare an operator functional if its launcher writes into an input, including recurrent state, cache, workspace, or output buffers. Prefer functional operators when practical.