Building an Out-of-Tree PyTorch Backend, End to End

A PyTorch backend is what makes model.to(device) and torch.compile(model) work on hardware PyTorch was not written for. It registers the device, manages the device’s memory, runs operators on it and compiles graphs for it. The vendor backends that do this are large, and their documentation describes the API they ended up with. It rarely says why the pieces have the shape they have, or in what order someone starting from scratch should build them.

This article walks through that process, from the first commit to a backend that compiles transformer and ResNet blocks, fuses them into generated CUDA kernels, calls cuBLAS and cuDNN where a library is the better choice, and checks every result against PyTorch. The example throughout is tessiq, a backend I built to work these questions out. It targets NVIDIA GPUs, has a CPU reference device beside them, and is small enough to read in full. The decisions it makes are the ones any backend of this kind faces, and most of them carry over to other hardware.

The article assumes you know PyTorch well and CUDA at the level of kernels, threads and memory, and that you have not built a backend before. The exact implementation of a device registration or a pass differs from one codebase to the next, but what it has to do is the same everywhere, and that is what the sections below describe. The numbers quoted come from one machine, an NVIDIA GeForce RTX 5080 with 84 multiprocessors, 64 MiB of L2 cache and 960 GB/s of DRAM bandwidth, running Linux with CUDA 12.8 and PyTorch 2.11. Later sections call it the test machine or the card.

Contents

1. What PyTorch asks of a backend
2. The shape of the system
3. An order to build it in
4. The repository and the build
5. Errors as values
6. The device runtime
7. A caching allocator
8. Registering the device with PyTorch
9. An IR of the compiler’s own
10. Passes and the pass manager
11. From torch.compile to the IR
12. Deciding what runs on the device
13. Fusion and code generation
14. Reductions, copies and library calls
15. Recipes and the executor
16. Shapes that arrive with the call
17. Custom operations
18. Precompiled operations
19. Measuring a backend
20. Testing a backend
21. What comes next
22. Porting to other hardware

What PyTorch asks of a backend

A backend meets PyTorch at two interfaces, and a complete backend implements both.

The first is the dispatcher. Every operator call in PyTorch, from torch.add to a convolution, is routed by dispatch keys. Each tensor carries a set of them, saying which backend holds its data and which cross-cutting features, such as autograd, apply to it. PyTorch reserves one key and one device type, both called PrivateUse1, for backends that live outside its source tree, so a new device needs no fork of PyTorch. Under that key the backend supplies everything that touches the device. That means memory allocation, copies to and from the host, a notion of the current device, and an implementation of every operator a program calls on the device’s tensors. This is the eager path, one operator at a time, as ordinary PyTorch code runs.

The second is the compiler interface of torch.compile. Dynamo captures a model as a graph while it runs, AOTAutograd normalises that graph and derives the backward pass, and the result goes to a function registered under a backend name. Inductor, PyTorch’s own compiler, sits at this interface. An out-of-tree device can plug its own code generation into Inductor, or bring a compiler of its own, a choice discussed below. This is the graph path.

The two paths depend on each other. A compiled graph nearly always contains operations the compiler will not take, and those have to run somewhere. Views and other cheap operations in the middle of a compiled graph run best through the eager path on the device. And a user who moves a tensor to the device and calls an operator on it, outside any compiled region, expects that to work.

Figure 1. The eager path and the graph path. The executor launches kernels where it has them, sends claimed nodes without a kernel through the eager path on the device, and runs unclaimed nodes on the host.

Four words in the figure recur through the whole article. A node is claimed when the compiler decides the device will run it, and unclaimed nodes stay on the host. A cluster is a group of neighbouring claimed nodes that run on the device together, so data crosses between host and device only at its edges. A recipe is everything the compiler decided about one graph, which node runs where, which kernels exist and which nodes they compute, recorded once so that a call does not decide it again. The executor walks the graph on every call and carries the recipe out. Each of them has a section of its own below.

PyTorch provides a great deal on both sides. The dispatcher, autograd, graph capture, functionalisation, the decomposition tables, and fake tensors (which carry shapes and dtypes without any data) all come with it. What a backend writes sits below and beside those pieces. On the eager side that is a device runtime, an allocator and the registrations that connect them to the dispatcher. On the graph side it is an importer, a compiler, kernels, and something that runs the compiled result on every call.

Inductor or a compiler of your own

The eager side has to be built either way. Before starting the graph side, decide whether it is a compiler of your own or a plug-in to Inductor, because the two lead to different projects.

Inductor lowers a graph to a loop-level representation, fuses and schedules it, and generates code through per-device classes. A new device can register its own with register_backend_for_device, passing the device name, a scheduling class that decides how fused loops become kernels, and a wrapper code generator that writes the host code calling them. A device on this path also registers a device interface with Dynamo, through register_interface_for_device, and its own operator overrides for code generation, through register_device_op_overrides. A device that can run Triton, or whose programming model is close to CUDA’s, gets fusion, scheduling, memory planning and dynamic shapes from Inductor this way, and its compiler work is concentrated in code generation.

A compiler of its own fits better when the hardware does not look like a GPU running loops. A dataflow or systolic accelerator, a device whose vendor already ships a graph compiler, or one that wants whole-graph decisions Inductor does not make, all fall here. In the last two cases the backend’s compiler is often mostly an importer and a partitioner that hand graphs to the vendor’s toolchain.

tessiq builds its own, from the importer to the generated CUDA, because the point of the project was to see every piece a graph compiler contains. For a GPU-like device the sections on the IR, passes and fusion describe work Inductor would do for you, and the rest applies either way. The runtime, the allocator, the registrations, the decision of what runs on the device, the library bindings, custom operations, measurement and testing are all needed whichever compiler sits in the middle.

The shape of the system

tessiq is a Python package over a C++ extension module, and most of the C++ does not depend on PyTorch at all.

Figure 2. The layers. Arrows point from a layer to what it depends on.

The core is four C++ libraries. The support library holds the error type and logging. The runtime holds the device interface, the CPU and CUDA backends behind it, and the caching allocator. The IR library holds the intermediate representation, the passes and memory planning. The code generation library holds the CUDA emitter, the wrapper around CUDA’s runtime compiler, and the cuBLAS and cuDNN bindings. These build with the standard library and, optionally, the CUDA toolkit.

Only the extension module includes PyTorch’s headers. It holds the dispatcher registrations, the allocator adapter PyTorch talks to, the importer that reads PyTorch’s graphs, the kernel objects the executor launches, and custom operations. Recent PyTorch headers need C++20, so the extension builds as C++20 while the core stays on C++17.

Keeping PyTorch out of the core has three consequences. The compiler can be developed and tested from text files on a machine with no GPU and no PyTorch, through a command-line driver. A change to a pass rebuilds a small library and leaves the extension alone. And everything PyTorch hands the compiler is translated at one boundary, the importer, so a pass never sees a Python object and the IR’s types are the compiler’s own.

The Python package is the glue. It registers the device’s Python-facing pieces, provides the function torch.compile calls, caches compiled results, runs them, and holds the profiler and the tools that explain a compilation. Code that handles PyTorch’s Python objects, such as graph modules, fake tensors and the compile options, lives there. Code that runs once per element, once per node or once per pass lives in C++.

One convention runs through the rest of the article. tessiq’s device index 0 is a CPU reference device whose “device memory” is ordinary host memory, and indices 1 and up are GPUs. In Python they have names, tessiq.cpu, tessiq.gpu and tessiq.gpu_at(n), so user code does not depend on the numbering. The device runtime explains why the reference device exists and why it comes first.

An order to build it in

A backend has many parts, and the order they are built in decides how long the project goes without anything that works end to end. I built tessiq in stages, and each stage ended with something that ran and was tested.

1. Build, packaging and CI. The repository layout, the CMake build, the Python package, the test runners and the CI gates, before any feature.
2. A device runtime and eager execution on a CPU reference device. Allocation, copies, the caching allocator, the dispatcher registrations, and a CPU fallback for every operator without an implementation of its own. At the end of this stage tensors can live on the device and any PyTorch operator runs on them, slowly.
3. A CUDA device. A second implementation of the same runtime interface. The allocator and the copy operators do not change.
4. The IR and the pass infrastructure, with no PyTorch dependency. A text format, a verifier, a pass manager, a few general-purpose passes and a command-line driver.
5. Graph capture and partitioning. The torch.compile entry point, the importer, and the passes that decide what the backend claims. Unclaimed operations run on the host inside the compiled call, so every model runs from the first day, however little of it is claimed.
6. Shape inference, memory planning and a cache for compiled graphs.
7. Elementwise fusion and CUDA code generation, compiled while the program runs.
8. Operator coverage for real model blocks. A table of operations, reductions, matrix multiplies through cuBLAS, and layout handling, until a transformer encoder block and a ResNet residual block match PyTorch.
9. Training and dynamic shapes. Backward graphs through the same compiler, and graphs whose sizes are not known until the call.
10. Device-side copies, so that rearranging a tensor’s layout no longer goes through the host.
11. Measurement. A profiler and device timing, to find where a call’s time goes.
12. Convolution through cuDNN.
13. Custom operations, including rewrite passes that find a hand-written computation in a model and replace it with one call.
14. Precompiled operations, whose host code is built by nvcc, for kernels that need libraries such as CUTLASS and cuBLASLt.

Three rules produced this order. Correctness comes before speed, so the CPU fallback and the reference device come first, and every later stage replaces something slow that already worked. Each stage has to be testable on its own, which is why the IR exists before the importer and can be exercised without PyTorch. And a measurement stage comes before the performance work that follows it, the convolution binding and the hand-written operations, because the first guesses about where a compiled call spends its time are usually wrong.

The sections below follow roughly the same order, the build first, then the runtime and the dispatcher, then the compiler. Measuring and testing come last because they apply to every part. Tests were written with every stage, not after them.

The repository and the build

The first stage has no features, and it is still worth doing carefully, because every later stage inherits its conventions.

The C++ side is one CMake project with a handful of options. They turn the C++ tests, the Python extension and the CUDA backend on or off, and make warnings errors. The CUDA backend is on by default when the build finds a CUDA toolkit and off otherwise, and the rest of the code does not need to know which kind of build it is in. The Python wheel is built by scikit-build-core, which drives the same CMake project, and pybind11 provides the bindings. The package version is read from the CMake project, so the wheel, the native module and the build always agree.

The repository is laid out by what each part depends on.

• csrc/ holds the C++. The support, runtime, IR, pass and code generation libraries each have a directory, the code that includes PyTorch has its own, and so do the Python bindings and the custom operations.
• python/tessiq/ holds the Python package.
• tools/ holds the command-line driver for the compiler, the generator for the operation table, and a script that explains a compilation.
• tests/cpp/ holds GoogleTest suites and a directory of golden files for the passes, and tests/python/ holds the pytest suites.
• bench/ holds benchmark scripts, each of which states what it measures.
• ops.toml, at the root, is the table of operations the compiler knows.
• docs/ holds the architecture description and short design notes, each recording one decision and the reason for it.

CI enforces warnings as errors, clang-format, clang-tidy, ruff and mypy in strict mode, and builds and runs the C++ tests a second time under AddressSanitizer and UndefinedBehaviorSanitizer. The CI machines have no GPU, so CI runs the C++ tests and every Python test that does not need one. Tests that need a GPU carry a marker and run on a development machine. Testing a backend covers how much of the backend CI still exercises that way.

One build detail catches every project that uses self-registration. Several things in a backend register themselves when a library loads. Passes add themselves to a registry, dispatcher kernels register with PyTorch, and custom operations record their definitions. Each does it in the initialiser of a global object. A linker takes an object file from a static library only when something references a symbol in it, and nothing references a self-registering pass by name, so an ordinary link drops every such file and the registry comes up empty without an error. The fix is to link those libraries whole-archive, which keeps every object file, with a comment beside the link line saying why.

Errors as values

Before writing any code that touches hardware, decide how failures travel. In tessiq the core never throws. Every call that can fail returns a Result, which holds either a value or an error, and an error carries a code and a message. The codes are few.

• Invalid argument, for a malformed request, such as a device index that does not exist.
• Unsupported, for a well-formed request this target cannot serve. Callers treat it as a signal to fall back rather than to fail.
• Not found, for a lookup that found nothing.
• Resource exhausted, for running out of memory.
• External, for a failure in CUDA, NVRTC, the filesystem or anything else outside the backend.
• Internal, for a bug in the backend found by a check that reports rather than aborts, such as the verifier finding a corrupted graph.

Why not exceptions? The reason is particular to compilers. In a compiler, failure is often the ordinary answer. A pattern does not match, there is no kernel for a dtype, a cache lookup misses, and those answers decide what to fall back to. A function that returns a Result says in its signature that failure is expected, and the code lets a caller branch on “unsupported” without parsing a message. The type is marked so that discarding a returned Result is a compiler warning, which the build makes an error, and two small macros propagate an error in one line.

Exceptions appear only at the edges, where the other side expects them. The Python bindings turn each code into the matching Python exception, such as ValueError for an invalid argument and MemoryError for exhausted memory. Code that PyTorch calls through its own interfaces, such as the allocator or a dispatcher kernel, raises PyTorch’s own errors, because that is how PyTorch expects failure to be reported there.

The device runtime

The runtime is the part of a backend that touches hardware. It allocates device memory, copies bytes between the host and the device, and waits for the device to finish. tessiq puts all of that behind one C++ interface, Device, with two implementations. One is a reference that keeps its “device” memory on the host. The other drives a CUDA GPU.

What the interface contains

The interface is small on purpose. A device reports what kind of backend it is and which physical unit it drives, allocates and frees raw memory, copies in three directions, host to device, device to host and device to device, and waits for its queued work to finish. Those are eight pure virtual methods. A ninth reports free and total memory and has a default for backends that cannot answer. Freeing cannot fail, because it runs from destructors. Everything else that can fail returns a Result.

The interface deliberately has no streams, because each backend orders its own work, as described below. Nor can it launch a kernel. Kernels belong to the compiler, whose code generation layer talks to CUDA directly. The interface isolates memory and data movement from the hardware. Code generation has a seam of its own, the emitter, described in Fusion and code generation.

The payoff comes with the second backend. The caching allocator calls the interface to allocate and free without knowing what is behind it, and the copy operators call its three copy methods. When the CUDA backend was added, neither changed.

The interface is thread-compatible rather than thread-safe. The allocator serialises its own calls into a device, and copies and synchronisation reach a device without a lock, which is sound because the CUDA runtime is itself thread-safe. Any state a backend creates lazily, such as a stream made on first use, has to be protected inside that backend.

A reference device first

The first implementation keeps device memory on the host. Allocation is an aligned malloc, every copy is a memcpy, and synchronisation returns at once, because nothing runs asynchronously.

With it the whole backend runs on any machine. A model compiled for the reference device goes through allocation, copies, the dispatcher registrations, partitioning, the compiler and the executor, with every claimed operation running through the CPU fallback, and its results can be compared with PyTorch on the CPU. That is how most of the test suite exercises the backend, including in CI, where there is no GPU. It also means device 0 always exists, so no code has to handle a runtime with no devices.

Building it first tests the interface as well. Once the allocator, the copy operators and the dispatcher registrations work against a device that is really host memory, the CUDA backend only has to get its own eight methods right.

The CUDA backend

The CUDA implementation maps each method onto the CUDA runtime API. Allocation is cudaMalloc, copies are cudaMemcpyAsync on a stream, and synchronisation waits for the device. Running out of memory is reported with its own error code. Every other CUDA failure becomes an external error that carries both CUDA’s symbolic error name and its message, since the name is what a developer searches for and the message is what a user reads.

The current device needs more care than any of the calls. The CUDA runtime keeps a current device per thread, and every call that does not name a device acts on it. The backend keeps its own current device index per thread as well, which PyTorch sets through the device guard described in Registering the device with PyTorch. Nothing keeps the two in agreement. PyTorch’s own CUDA code can change the CUDA setting at any time, and with two GPUs a thread can be working on the second while CUDA still names the first. So each CUDA device object knows its own ordinal and selects it before every call into the driver. Selecting a device that is already current costs little.

The backend is also lazy. Constructing a device object touches nothing. The first call that needs the device selects it and creates its stream, and that first call is also when the driver creates the device’s context (the driver-side state for one process on one GPU). Creating a context takes long enough to notice and reserves memory on the card, so a design that counted the GPUs in a machine by constructing a device for each would initialise every one of them.

Device indices and ordinals

PyTorch identifies a device by a type and an index, and an out-of-tree backend gets exactly one type. tessiq uses the index to choose the backend as well as the unit.

Figure 3. A device index selects the backend and the unit. Index 0 always exists, and GPUs follow in the order the driver reports them.

Inside the CUDA backend the physical unit is called the ordinal, which is the number cudaSetDevice takes, and it is kept separate from the index users see. Index 1 is ordinal 0. Mixing the two up passes the driver a number one too high, which fails on a machine with one GPU and silently uses the wrong card on a machine with two.

The simpler design is a process-wide switch that sends every device tensor either to the CPU or to the GPU. Choosing per tensor costs almost nothing more and lets one process use both backends, which the test suite relies on. It compiles the same model for the reference device and for a GPU in one process and checks both against PyTorch.

The devices live in a registry that is built on first use. It creates the reference device, asks a discovery function for the GPUs, and gives every device its own caching allocator, so memory statistics are per device. Discovery has to work on a machine with no GPU, with no driver, or with a driver older than the CUDA runtime the backend was built against. All three mean there is no CUDA to offer, so discovery returns an empty list, logs the reason, and the runtime carries on with the reference device. A build without the CUDA toolkit compiles discovery to a function that returns an empty list, so nothing above the registry knows which kind of build it is in.

Ordering work without waiting for it

A stream is a queue of GPU work that runs in order. Each CUDA device creates one stream and queues every copy on it. Kernels, which the compiler launches, go to a different stream, the legacy default stream, which is the one a CUDA call uses when it names no stream at all.

The split works because of how the legacy default stream behaves. A stream created with cudaStreamCreate is a blocking stream, and the legacy default stream synchronises with every blocking stream on the device. Work queued on it waits for work queued earlier on the blocking streams, and they wait for it in turn. A kernel therefore runs after the copy that feeds it, the copy of its result runs after the kernel, and nothing in the backend has to order a copy against a kernel explicitly.

Figure 4. How a copy in, a kernel and a copy out are ordered across two streams. CUDA inserts the waits because one of the two streams is the legacy default stream.

Synchronising the device waits for both streams, since either may hold the most recent work. A copy the caller asked to be blocking synchronises before it returns, so the host sees the bytes it asked for.

The price is overlap. A copy cannot run while a kernel does. For a compiled call that costs little, because a call moves its inputs at the start and its result at the end, with no independent work in between. A backend that wants to pipeline transfers with compute needs separate non-blocking streams and events to order them, and its allocator then has to track which stream each block was used on, as the allocator section explains.

Copies between devices

Every copy between two device indices is staged through the host. Between the reference device and a GPU there is nothing else to do, since the reference’s memory is host memory. Between two GPUs, staging is a choice. Within one process CUDA gives all GPUs one virtual address space, and cudaMemcpyPeerAsync copies between them directly, over the interconnect where the hardware has one. A backend meant for machines with several GPUs should use it. Staging is correct everywhere and was enough for a machine with one GPU.

What a port changes

The device implementation is the core of a port to other hardware. The CUDA one maps onto any runtime that has an allocator, a memory copy and a way to wait. The caching allocator and the copy operators come with it unchanged. The rest of a port, most of it in code generation, is collected in Porting to other hardware.

A caching allocator

A training step allocates and frees device memory all the time, for activations, gradients, optimizer temporaries and the result of nearly every operator. If each of those went to the driver, the backend would spend its time waiting. cudaMalloc and cudaFree are slow next to a kernel launch, and cudaFree synchronises the device. So every production backend puts a caching allocator between PyTorch and the driver. It takes memory from the driver once, hands it out, and keeps freed blocks for the next request instead of giving them back.

How PyTorch asks for memory

PyTorch asks through c10’s allocator interface, whose central method takes a byte count and nothing else. There is no device argument. The allocator registered for a device type allocates on the thread’s current device, as CUDA’s does. What it returns is a DataPtr, an owning pointer tagged with the device the bytes are on, together with the function to call when the last owner lets go. A storage (the untyped block of bytes under every tensor) holds one.

tessiq’s adapter reads the thread’s current device index, asks that device’s caching allocator for the bytes, and wraps them in a DataPtr. There is one caching allocator per device, and one adapter instance bound to each device index as well, for code that allocates on a known device, such as a factory operator that was given an explicit device. The adapter derives from c10’s device allocator class instead of the plain allocator, which is what lets torch.accelerator.memory_allocated() and its relatives read the backend’s counters. It is one of the few places where an error becomes an exception, because c10 reports allocation failure by throwing.

A deleter without context

The other end of the interface has the same gap in reverse. The deleter in a DataPtr is a plain function pointer, and c10’s contract for allocators that expose one requires the context it receives to be the data pointer itself. So when a block is freed, the deleter gets an address and nothing else, no device and no allocator. With a pool per device it has to work out which pool the address came from, and handing a block to the wrong device’s allocator would give it to one that never handed it out.

The adapter therefore records the owning device of every address when it allocates, in a map from address to device index behind a mutex, and looks the address up and removes it when the block is freed. That costs a hash lookup per allocation and per free, and a backend with a single device can skip it. The same map answers PyTorch’s question of which device a raw pointer belongs to.

Size classes

Every request is rounded up to a size class before anything else happens. Requests under 1 MiB round up to a multiple of 512 bytes, and larger ones to a multiple of 2 MiB. Freed blocks go into a free list keyed by their rounded size, and a request is served from the list for its exact class or, on a miss, from the driver.

Rounding is what makes reuse happen at all. Without it the free list fills with blocks of sizes that never come back, and the cache holds memory without reusing it. Two granularities are needed because one is wrong at one end or the other. Rounding a 100 MB tensor to 512 bytes gives a class that almost no other request will share, and rounding a 4 KB tensor to 2 MiB wastes nearly all of it. Rounding small requests finely and large ones coarsely means activations whose sizes differ slightly from one iteration to the next still land in the same class. The constants match PyTorch’s own CUDA allocator, which also rounds small requests to 512 bytes and treats 1 MiB as the boundary between small and large.

Figure 5. The life of a block in the caching allocator.

Running out of memory

When a driver allocation fails, the allocator is often holding free memory of its own, in blocks cached for classes nobody asks for any more. The caller cannot know that. So before reporting failure the allocator returns every cached block to the driver and tries once more. If nothing was cached there is nothing to retry with, and the error goes straight up. If the retry fails as well, the error says how many bytes were released on the way, which tells whoever reads it whether the device was really full or the cache had grown. The retry does not look at why the first attempt failed. Any failure is worth one more attempt with the cache empty.

Freeing never returns memory to the driver. A freed block stays in its free list until a request of its class takes it, or until the user empties the cache.

Counters a test can check

The allocator counts what it does, including the bytes live tensors hold, the bytes taken from the driver and not given back, cache hits and misses, retries and failures. These counters are part of the design rather than a debugging aid, because they are how a test proves the cache behaves. Allocating a 4000-byte tensor, freeing it and allocating it again should leave 4096 bytes allocated, the same 4096 reserved, and one hit. The counters also show a workload’s cost. In a long-running program, a gap between reserved and allocated memory that keeps growing means the cache is collecting classes that are never requested again. And a compiler test can read them to prove that work ran on the device at all, as Testing a backend describes.

Alignment

The device interface promises that every pointer it returns is aligned to 512 bytes, and nothing above it realigns. Code that depends on alignment can rely on that for the start of every storage, though not for a view at an offset inside one. The CPU reference meets the promise with an aligned allocation. For CUDA, the documented guarantee of cudaMalloc is 256 bytes. The 512-byte promise held on every allocation measured, because the driver packs small allocations at 512-byte steps and gives large ones 2 MiB alignment or more. Because the promise is stronger than the documented guarantee, the interface states it next to the allocation method, together with the observation it rests on.

Why no stream tracking is needed

A block freed on the host may still be in use on the device. Freeing happens when the last tensor that refers to a storage is destroyed, which can be right after the kernel that reads it was queued and before that kernel has run. If the allocator hands the block to a new tensor and a new kernel writes it, the old kernel reads garbage.

PyTorch’s CUDA allocator handles this with a pool per stream and a record_stream call that marks a tensor as used on another stream, so its block is not reused until that stream has caught up. tessiq tracks no streams, because all of its device work is queued on streams that are ordered against each other, as described in Ordering work without waiting for it. Any work that reuses a freed block is queued after the work that used it, and so it runs after it. The adapter’s hook for recording a stream is empty, and it would have to become real the day the backend used independent streams.

Pinned host memory

Copies between the host and a GPU are fastest, and can be asynchronous, when the host memory is page-locked, since the driver can then move it by DMA without staging it first. PyTorch asks the accelerator’s hooks object for an allocator to use for tensor.pin_memory() and for a data loader that pins its batches. tessiq’s pinned allocator uses cudaHostAlloc when the machine has a GPU and ordinary host memory when it does not, since then every device is host memory and pinning buys nothing. Freeing a page-locked block first synchronises the GPUs, because a copy started without blocking may still be reading from it.

What the simple design gives up

The whole allocator is written against the device interface. Compared with PyTorch’s CUDA allocator it gives up four things.

• Every block is its own driver allocation, where PyTorch carves blocks out of larger segments.
• A request is served only from a cached block of exactly its class. A larger free block is never split to serve it, and free neighbours are never merged.
• Large requests round to 2 MiB, so a request for 1 MiB and one byte takes 2 MiB.
• There is one pool per device, where PyTorch keeps one per stream.

The first three cost driver calls and memory on workloads whose sizes keep changing, and the counters show that cost as the gap between reserved and allocated memory. The fourth costs nothing while all device work is ordered as described above. A backend that expects widely varying shapes, or wants transfers to overlap with compute, should start from PyTorch’s design instead.

Registering the device with PyTorch

With a runtime and an allocator in place, PyTorch has to be told the device exists. torch.randn(4, device="tessiq") has to allocate memory the backend owns, x + x has to produce the right numbers, and .cpu() has to bring them back. All of that goes through the dispatcher, the component that decides which implementation of an operator runs for a given call.

How the dispatcher routes a call

For every operator the dispatcher keeps a table from dispatch key to kernel. A call combines the key sets of its arguments, takes the highest-priority key, and runs the kernel registered for it. The autograd key ranks above the backend keys, so its kernel runs first, records what the backward pass will need when an input requires a gradient, and then passes the call on to the backend.

Two kinds of kernel are shared by every backend, and they decide much of what a new one has to provide.

• A CompositeImplicitAutograd kernel implements an operator by calling other operators, and it runs before any backend kernel is consulted. softmax is one. It calls _softmax, and only that reaches a backend.
• A CompositeExplicitAutograd kernel is a default that any backend without its own kernel receives. permute is one. It computes new strides and calls as_strided.

Counted on PyTorch 2.11, the dispatcher holds 3203 ATen operator overloads. An overload is one signature, so add.Tensor, which adds two tensors, and add.Scalar, which adds a tensor and a number, are two. 838 of the 3203 have a CompositeImplicitAutograd kernel, which leaves 2365 that a backend can receive. No backend implements all of them, and the design has to cope with that from the first day.

Figure 6. How a call on device tensors is routed once the autograd kernel has run. CompositeImplicitAutograd operators have already been broken into other operators before this point.

What a device registers

Seven registrations make the key lead somewhere. Three are C++ objects created while the extension library loads, and four are Python calls made when the package is imported.

Table 1. What a PrivateUse1 backend registers with PyTorch, and the API each registration goes through.

The C++ registrations are global objects whose constructors do the registering. They run during static initialisation, when the dynamic loader maps the extension into the process and before any of its functions are called, so loading the library is enough. Their relative order is not under the program’s control, so none of them may depend on another.

The device module is the registration most often missed. When a call names the device, PyTorch looks for a module at torch.tessiq, and during lazy initialisation and seeding it calls into that module for the device count, the current device and the state of the random number generator. Nothing provides the module by default, so the backend builds one and registers it with torch._register_device_module. It carries what PyTorch looks for, such as is_available, device_count, current_device, set_device, synchronize, manual_seed and get_rng_state. In tessiq the seeding functions accept the seed and do nothing. Random operators on a device tensor run on the host through the fallback and draw from PyTorch’s default CPU generator, which torch.manual_seed has already seeded. torch.manual_seed also calls every device module’s seeding function, so one that called back into it would recurse.

The device guard

A device guard makes a device current for a scope and restores the previous one when the scope ends. PyTorch’s device-generic code builds its guards on the backend’s guard implementation, which gets and sets the current device, counts devices, and reports streams. tessiq’s reports only a default stream, since it orders all the work on a device as if it were on one. The users are c10’s generic guards, the autograd engine’s per-device worker threads, and the device functions in torch.accelerator. The current device itself is a thread-local index in the runtime, and the guard implementation reads and writes it.

The guard is not applied around a backend’s own kernels. For CUDA and the other backends built into PyTorch, PyTorch’s code generator writes a device guard into the registration of each operator that needs one. Kernels registered by hand get none. A kernel registered under PrivateUse1 and called with a tensor on device 1 still sees whatever device was current before the call.

So every kernel that allocates, or otherwise depends on the device, sets the device itself, with a small scoped guard from the runtime, and passes the allocator for that index explicitly, trusting nothing about which device is current. One convention needs care here. An absent device, or one with a negative index, means the current device. That is what PyTorch passes for torch.empty(4, device="tessiq"), so a kernel that treated it as an error would break the most common call there is.

The CPU fallback

Every operator that has no kernel of its own under the key, and no CompositeExplicitAutograd default, goes to one registration, a fallback for the whole key. One function can serve thousands of operators because it is boxed. An unboxed kernel has the operator’s own C++ signature, so it serves only that operator. A boxed kernel has one signature for everything. Its arguments arrive on a stack of IValue (a tagged union that can hold a tensor, an integer, a list or anything else an operator takes), and its results go back on the same stack.

What makes a generic handler possible is the schema, the declared signature every operator has. It gives the operator’s name, the types of its arguments and results, and which arguments it writes in place. PyTorch ships a CPU fallback that reads the schema to find the tensors among the arguments, copies them to the host, and calls the operator again with the CPU key. Results come back to the device with .to(), and any argument the operator wrote in place is copied back into the original device tensor. tessiq’s fallback logs the operator and calls PyTorch’s.

Two transfers per operator is slow, and it buys two things. Most of ATen works on the device from the first day, so kernels can replace the fallback one operator at a time without anything breaking in between. And the fallback never becomes dead code, because the operators a backend does not implement have to run somewhere. In compiled graphs the compiler keeps unclaimed operations on the host in the first place, so the fallback’s cost falls mostly on eager code.

The fallback is not right for every operator, and the exceptions decide which operators must be native. It copies, so an operator whose result must share memory with its input loses that property, and every view a backend supports has to be registered natively. And an operator with a CompositeExplicitAutograd default never reaches the fallback, even when that default only raises.

Operators that have to be native

A small set of operators is registered under the key because the fallback would give them the wrong behaviour.

• Factories. empty.memory_format and empty_strided create a tensor rather than transform one, so there is nothing to copy to the host. They set the device with the scoped guard and call PyTorch’s generic empty-tensor construction with the device’s allocator and the key.
• Copies. _copy_from implements copy_ between a device tensor and anything else, which is what .to() and .cpu() come down to. Its fast path moves a whole storage in one transfer, which is correct when source and destination have the same dtype and layout and each covers its whole storage. Otherwise the copy stages through the host and lets PyTorch’s CPU kernels convert the layout or the dtype, so the backend needs no second copy of PyTorch’s striding and type-promotion rules. _copy_from_and_resize is the variant the fallback uses to write back an argument that was modified in place.
• resize_. Resizing changes the storage a tensor points at, in place. Through the fallback it would resize a host copy, and writing that back would need a resize of the device tensor, which is the operation being implemented.
• set_. torch.load rebuilds each tensor by pointing an empty tensor at a storage. Through the fallback, a host copy would be pointed at the storage and the real tensor left empty.
• Views. as_strided, view and _reshape_alias have kernels only for PyTorch’s own backends, so on a new device they would reach the fallback and return copies. ATen’s implementations of all three touch only metadata and work on any device, so the backend registers them as they are. permute, transpose, slicing and most other views are CompositeExplicitAutograd kernels that call as_strided, so these three registrations cover them.
• clone. clone, which .contiguous() also reaches, runs a device-side copy kernel on a GPU that packs any strided layout into a new tensor, described in Reductions, copies and library calls.
• convolution_overrideable. PyTorch sends a convolution on a device it has no convolution for to this operator, which exists for backends to implement. Its CompositeExplicitAutograd default raises, and because a default exists, the fallback never sees it. tessiq implements the eager version by running PyTorch’s convolution on the host, and compiled graphs call cuDNN. A backend that trains convolutions eagerly registers the backward counterpart, convolution_backward_overrideable, in the same way.

One accelerator per process

PyTorch has a notion of the accelerator, one device type per process, which torch.accelerator.current_accelerator() returns. Registering a PrivateUse1 backend makes it that accelerator. After the package is imported the call returns the new device where it returned CUDA before.

The effect reaches CUDA tensors. The autograd engine asks the accelerator for the stream each node of the backward graph ran on, and a CUDA tensor is not on the new device, so a backward pass over CUDA tensors fails inside the engine. Importing the package is enough to cause it. A process can use the new backend or CUDA autograd, not both. Forward passes on CUDA are unaffected, which is why tessiq’s benchmarks can compare inference against CUDA and Inductor in the same process, while its tests compare against PyTorch on the CPU. tessiq records the constraint in one of its design notes.

An IR of the compiler’s own

torch.compile hands a backend a graph, so why build a second one? tessiq does. It imports every graph PyTorch gives it into an intermediate representation (an IR) of its own, and every compiler decision after that point is made on the IR. The IR depends on neither PyTorch nor CUDA and has no notion of which device anything will run on, which is what lets the compiler be developed and tested from text files.

Why not rewrite FX directly

FX is PyTorch’s graph format. Dynamo produces an FX graph, and AOTAutograd produces the one a backend receives. It is a capable structure. Its nodes know their users, and it has a lint that checks every argument is defined before it is used. Four of its properties make it the wrong thing for a compiler like tessiq’s to rewrite.

• It is a graph of Python objects, which passes written in C++, and a command-line tool that runs them without Python, cannot use.
• A value’s type is not part of the value. It lives in an optional metadata dictionary, when something has put it there.
• Arguments can be any Python object, from numbers and lists to dtypes and devices, so a pass that compares two nodes has to understand all of them.
• It has no text form that parses back into a graph, so a test of a pass cannot be a file with an input and an expected output.

tessiq’s IR is built around the opposites, a C++ structure with no Python in it, a static type on every value, a closed set of attribute types, and a text form that round-trips. It keeps the two good properties FX has, use lists and a verifier.

Values, nodes and use lists

A graph holds its parameters, an ordered list of nodes, and its results. A node has an operation name, ordered operands, ordered results and attributes. A value is either a graph parameter or one result of one node, and it is assigned exactly once. That is static single assignment, SSA. Where ordinary code might reassign a variable, SSA gives each assignment its own value, so every use of a value names exactly one definition.

SSA turns the relationship between a definition and its uses into a pointer rather than a question about control flow. Every value knows the node that produced it and the nodes that read it, its use-def chain, and most rewrites reduce to a question about one side or the other. Dead code elimination asks whether a value has readers. Common subexpression elimination moves the readers of one value to another. Constant folding asks whether an operand’s producer is a constant.

The use lists are maintained, never recomputed. Every mutation goes through a method on the graph or the node that updates both sides, whether it sets one operand, moves every use of one value to another, erases a node, or inserts one. Erasing refuses to remove a node whose results are still read, and names the readers in its error, so a value always outlives its uses. Inserting places the new node before an existing one, because a replacement has to be defined before its first use.

Figure 7. The parts of the IR, and the two tools that work on whole graphs.

Types that carry layout

Every value has a tensor type made of an element type, a shape and a layout.

The IR carries seven element types, fewer than PyTorch has. They are float16, bfloat16, float32, float64, boolean, and 32-bit and 64-bit integers. A graph that uses any other type is not imported and runs uncompiled.

A shape is a list of extents, and an extent can be unknown until the graph runs. The text form writes an unknown extent as a question mark. A float32 matrix with 16 columns and a number of rows that arrives with the input has a known and an unknown extent. A rank-0 shape is a scalar, distinct from a one-element vector. The IR could represent unknown extents from its first version, long before the compiler could compile a graph that had them, so supporting them later needed no change to the type system.

The layout is the stride of each axis, in elements. It is absent for a packed row-major tensor, which is almost every value, and present otherwise. The transpose of a packed 32 by 32 matrix, for example, has the same bytes with a stride of 1 on its first axis and 32 on its second. Layout is part of the type because several compiler decisions depend on it. Generated kernels read their inputs at a flat index, which is correct only for packed data, so whether a fused loop can read an operand directly depends on its strides. A matrix multiply’s transpose flags for cuBLAS are read from the strides. And where a claimed kernel reads a strided value, a pass has to record that a packed copy is needed first.

Operations and attributes

An operation is a name, not a member of a fixed enumeration. An imported graph carries the name PyTorch uses, such as aten.add.Tensor, a hand-written test can use a short name such as add, and a pass can introduce names of its own. What an operation means to the compiler, whether it is elementwise, what it computes, and whether the backend runs it, comes from a table keyed by name, described in Deciding what runs on the device. Keeping the IR open like this means importing a graph never fails because an operation is unfamiliar. An unfamiliar operation is simply one the backend does not claim.

Tensor-valued arguments are operands, which are edges in the graph. Every other argument is an attribute, and attributes are a small closed set of types, an integer, a float, a boolean, a string, a list of integers, or none. None is needed because graphs pass it, for example where a convolution has no bias, and dropping it would renumber every argument after it. Positional arguments are keyed by their position and keyword arguments by their name. Anything richer than this set is usually asking to be an operand instead, and the narrow set keeps the printer, the parser and every pass that compares attributes simple.

Passes record their decisions as attributes too, under a prefix of their own, marking whether a node is supported, which cluster it belongs to and which fusion group. Every decision is then visible in the printed IR, which is how a test checks it. Attributes are stored sorted by name, so a printed graph is the same on every run.

A text form and golden files

The IR prints and parses, and the reason is testing. A pass is a function from a graph to a graph, and with a text form a test of a pass is a file. Each golden file names a pipeline of passes, gives an input graph, then a separator, then the exact expected output. The test runner parses the input, runs the pipeline, prints the result and compares text, so a failure is a diff and adding a case is adding a file. tessiq has about fifty, and nearly half of them test a case where a pass must decline or leave something alone. Those are the cases that catch a matching condition widened too far.

The verifier

The verifier walks a graph and checks the invariants every pass is entitled to assume.

• Every operand is defined before the node that uses it, in list order.
• Every operand’s use list names the node that reads it.
• Every result points back at the node that produced it.
• No value name is used twice, and every graph result is defined in the graph.

The second is the one a careless pass breaks. A pass that assigned an operand directly instead of going through the method that maintains use lists would leave the new operand’s list without its reader, and a later pass asking who reads that value would get the wrong answer. So the verifier’s error names the likely cause along with the symptom. In debug builds the pass manager runs the verifier after every pass and reports the pass by name, so a pass that corrupts the graph is reported at the moment it does so rather than as a wrong result several stages later. The parser applies the same checks to hand-written input.

A command-line driver

tessiq-opt reads IR as text, runs a pipeline of passes named on its command line, and prints the result. It can also list the registered passes, verify a file without printing it, print the memory plan for a graph, and print the CUDA source the code generator would write for each fusion group. It links the IR, the passes and the code generator, and runs without PyTorch or a GPU.

It also sets the workflow for compiler bugs. When a pass misbehaves on a real model, dump the model’s IR, cut it down to the smallest graph that still misbehaves, and save that as a golden file, where it stays as a test.

Passes and the pass manager

A compiler is a sequence of transformations over its IR, and each transformation is a pass. This section covers how a pass is written, registered, combined with others into a pipeline, and tested.

The pass interface

A pass has a name, a one-line description, and a method that runs it on a graph and returns either whether it changed the graph or an error saying why it could not proceed. The name is how a pipeline selects it, and the description is what the command-line driver lists.

Each pass also documents three things in a comment above it. What it assumes on entry, what it guarantees on exit, and why its rewrite does not change what the graph computes. The first two decide where the pass can sit in a pipeline. The third, the legality argument, is the part nobody can recover from reading the code later. Dead code elimination’s is short, and it carries a condition. Removing a node nobody reads cannot change what the graph computes, because no operation in the IR has side effects. That holds because AOTAutograd functionalises every graph before the backend sees it, so no node writes to its inputs. The day an operation with a side effect enters the IR, the pass becomes wrong in a way no test over pure graphs would catch, and it will need an effects marker to check. The comment keeps that condition next to the code that depends on it.

Registration by name

One macro registers a pass. It defines a global object whose initialiser adds a factory, a function that builds a fresh instance of the pass, to a registry keyed by the pass’s name. The initialiser runs when the library loads, so writing the pass and the macro in one file is the whole registration. The pass then appears in the driver’s list, in the list Python can read, and in any pipeline that names it. This is the self-registration that needs the whole-archive link described in The repository and the build.

Pipelines

A pipeline is a comma-separated string of pass names. The pass manager turns the names into instances through the registry and runs them in order, recording for each whether it changed the graph and how long it took. In debug builds it runs the verifier after each one. A misspelt name fails with the list of passes that exist.

Because a pipeline is a string, it can come from a command line, a test file or a function argument, and nobody edits a central list to add a pass. The pipeline a torch.compile call runs is fixed in code, because what the backend claims is a property of the backend, not a choice made per graph. Custom operations add passes to the front of it, as described in Custom operations.

The changed flag has to be true

The boolean a pass returns looks like bookkeeping, and code that drives passes relies on it. A typical consumer is a loop that runs a pipeline until no pass reports a change, which is how a compiler reaches a fixed point when one pass creates work for another. If any pass reports a change it did not make, that loop never ends. tessiq’s own pipelines run each pass once, but the flag is part of the pass’s contract, and a pipeline added later should be able to loop on it.

So a test runs every registered pass twice over the same small graph and requires the second run to report no change. A transformation that gives the same result applied twice as applied once is idempotent. The small graph gives only some of the passes real work, so for the others the test proves little, and their golden files carry their coverage. It matters most where two passes split one logical transformation between them. Common subexpression elimination moves the readers of a duplicate node to the original and leaves the duplicate, now unused, for dead code elimination to remove. When it runs again, the duplicate is still there and is found again, so it has to report a change only when a use actually moved, not every time it finds a duplicate.

General-purpose passes

The IR ships with five standard passes. Dead code elimination walks the graph backwards, so removing a consumer makes its producer dead within the same sweep. Common subexpression elimination counts attributes as part of a node’s identity, since two nodes with the same operation and operands can differ in a constant. Constant folding inserts each folded constant before the node it replaces, so a definition still comes before its uses.

Algebraic simplification shows the standard of legality the rest of the compiler follows. It replaces a node with one of its own operands when an identity says they are equal, such as multiplying by one, negating twice or applying ReLU twice. An identity has to hold for every value, including both zeros, both infinities and every NaN. Adding zero is exact for integers and not for floats, because IEEE 754 has two zeros and -0.0 + 0.0 is +0.0, so replacing x + 0.0 with x changes the result when x is negative zero, and the lost sign can come back later as an infinity of the wrong sign after a division. The pass removes an integer addition of zero and keeps a floating-point one.

Shape inference derives each node’s type from its operands with a rule per category of operation, leaves operations without a rule alone, and merges what it derives with what the graph records, extent by extent. When an unknown extent meets a known one in a broadcast, the result stays unknown. PyTorch would type it with the known extent, because it installs a guard that recompiles the graph if that assumption fails. tessiq has no such guard, so its kernels check the real sizes at launch instead.

None of these passes runs when a model is compiled. The importer takes every type from the metadata PyTorch recorded while tracing, and the compile pipeline consists of the placement and fusion passes. The general passes serve hand-written IR and the driver, and shape inference serves the tests, which run it over imported graphs and require its rules to agree with the types PyTorch computed.

From torch.compile to the IR

torch.compile(model, backend="tessiq") is the entry point to everything the compiler does. Between that call and the first pass over the IR, PyTorch captures the model and normalises the graph, the string "tessiq" is resolved to a function, and that function imports PyTorch’s graph into the IR.

What Dynamo captures

Dynamo intercepts Python bytecode as a function runs and records the tensor operations it performs into an FX graph. The graph is faithful to what the program did, which is exactly what makes it hard to compile. It contains calls to modules and methods, not only operators. It can modify its own inputs in place, because models use in-place operations. And it draws on the whole of ATen, thousands of operator overloads, because that is what users write. Mutation is the worst of the three for a compiler. If one operation writes a tensor another reads, their order is fixed, and reordering, grouping and fusing operations is most of what a compiler does.

AOTAutograd normalises the graph

AOTAutograd sits between Dynamo and a backend and deals with all three. It functionalises the graph, rewriting each in-place operation as one that returns a new value, so no node writes to its inputs. It decomposes the graph into core ATen, a much smaller set of operators, by applying a table of rewrite rules. And it traces the backward pass ahead of time as a second graph in the same operator set, which is where the “AOT” in its name comes from.

A backend opts into all of this by handing its compiler to AOTAutograd instead of taking Dynamo’s graph directly. tessiq’s entry function receives Dynamo’s graph, its example inputs and the options the user passed. It builds a forward compiler, a callback AOTAutograd calls with each graph it produces, and hands that callback to AOTAutograd together with the decomposition table. The callback runs the whole compiler on the graph, import, passes and kernels, and returns a function that runs the result. That function uses AOTAutograd’s boxed calling convention, a different sense of “boxed” from the dispatcher’s. It takes its arguments as one list that AOTAutograd can clear during the call. Separate arguments would stay alive in the caller’s frame until the call returned. In training those arguments include every activation saved for the backward pass. A graph the compiler cannot represent, such as one with a dtype the IR does not carry, runs uncompiled, and the compile still succeeds.

One option chooses the device that claimed operations run on, and defaults to the CPU reference. Another enables pattern rewrites for custom operations. Anything else is rejected when the backend first runs, which is the first call of the compiled model, not the call to torch.compile, since torch.compile itself only records the options.

Decompositions do most of the coverage work

A decomposition rewrites one operation as a composition of simpler ones while AOTAutograd traces. SiLU is a good example. The backend has no SiLU kernel. Without decompositions a graph that calls it is one node the backend does not support. With them, SiLU arrives as x * sigmoid(x), both of which the backend runs, and the graph goes from unsupported to fully claimed without anyone writing a SiLU kernel.

tessiq’s table starts from PyTorch’s core ATen decompositions and adds a few. Core ATen keeps softmax, layer norm and batch norm whole, because most backends have fused kernels for them. A backend without those kernels, but with every piece they decompose into, does better to decompose them. A softmax becomes a maximum, a subtraction, an exponential, a sum and a division. A layer norm becomes a mean, a variance and an elementwise pass. An inference-mode batch norm becomes a subtraction, a reciprocal square root, a scale and a shift.

Decomposing can go too far. PyTorch’s own decomposition of var_mean produces operations from prims, a lower operator set with its own names and semantics that tessiq does not cover. Descending one step too far turns a graph that was two operations short of complete into one that is short a whole namespace. So tessiq replaces that entry with its own, written in ATen as a mean, a squared difference, a sum and a division. The rule is to decompose until the pieces are operations the backend has, and stop there.

AOTAutograd modifies the dictionary it is given, so the table is built fresh for every compile. And a custom operation that takes over a PyTorch operation needs that operation to reach the compiler whole, so its decomposition has to step aside, as described in Custom operations.

Resolving the backend name

torch.compile receives a string, and nothing in the process has necessarily imported the backend when Dynamo reads it. There are two ways to make the name resolve. A backend can call Dynamo’s register_backend when its package is imported, which means the user has to import the package before compiling. Or the package can declare an entry point in the torch_dynamo_backends group of its metadata, naming its entry function. When Dynamo does not find a name in its registry, it scans installed packages for entry points in that group, imports the module named, and registers the function it finds. An installed package should use the entry point, and only the entry point. Dynamo registers what it loads through register_backend, which asserts that the name is not already taken, so a package that also registered itself on import would make the compile abort.

Resolving the name imports the package, and importing it registers the device. A plain torch.compile(model, backend="tessiq") is therefore enough to make the backend the process’s accelerator, with the consequences described in One accelerator per process.

Importing FX into the IR

After AOTAutograd an FX graph has three kinds of node, a placeholder for each input, a function call for each operation, and one output. The importer turns placeholders into graph parameters, function calls into IR nodes, and the output into the graph’s results. Method calls, attribute reads and module calls should no longer exist at this point, and meeting one is an error that says the graph did not come through AOTAutograd.

Types come from fake tensors. The IR requires a static type on every value, and FX has none of its own. What it has is metadata. While tracing, AOTAutograd runs the graph on fake tensors, which carry a shape, a dtype, strides and a device but no data, and it leaves each node’s fake result on the node. The importer reads the dtype, the shape and the layout from there, so it never has to work out what a permute or a slice does to strides, because PyTorch already did. A node without that metadata is refused rather than guessed at, because a wrong type raises no error at import and is then trusted by every pass that reads it.

Tuples become multiple results. A few operations return several tensors, and FX reads each with a getitem on the tuple. The IR has no tuples, since a node simply has several results. So a getitem is not imported as a node. It names one of them, and that result takes the name of the getitem that reads it, because the executor looks tensors up by FX node name.

Arguments split in two. Tensor-valued arguments become operands and everything else becomes an attribute. Python scalars, lists of integers and None have IR spellings. PyTorch enumerations the compiler never interprets, such as a dtype or a memory format, are carried as labels. Anything else is refused.

Symbolic sizes become values. A graph Dynamo traced for more than one shape passes its sizes as graph inputs and computes with them. The importer types those inputs as rank-0 integers and gives unknown extents the question mark from the IR’s shape type. Shapes that arrive with the call covers what the compiler does with them.

Figure 8. Both graphs AOTAutograd produces go through the same compiler callback. The backward graph is traced with the forward one and compiled only when a backward pass first runs.

The backward graph

AOTAutograd produces a backward graph for any model whose parameters require gradients, when grad mode is on, and the backend needs no code of its own to produce it. When no backward compiler is given, AOTAutograd uses the forward compiler for the backward graph as well, so both go through the same importer, passes and kernels. The backward graph is made of the same core ATen operations as the forward one. Supporting training therefore needs no derivative rules in the backend, which is the largest single saving AOTAutograd provides.

It also means an inference test suite exercises only half of the path. AOTAutograd traces the backward graph during the first forward call but hands it to the backend only the first time something calls .backward(). Training has to be tested on its own.

A backward graph returns one value per forward input, a gradient, or None for an input that takes no gradient, such as a frozen bias. The autograd engine matches the returned values to the forward inputs by position. A compiled graph that dropped the None entries and returned a shorter tuple would not crash. It would put one parameter’s gradient on another, and the model would train to something else with nothing to notice it. The IR records only computed values, so it cannot say how many positions return nothing.

The resolution is that the shape of what a caller gets back is part of the calling convention, not of the graph’s content. It is read from the FX graph, where the None positions are visible, stored beside the compiled graph as an output spec, and used by the executor to rebuild the returned tuple on every call. The executor pairs the spec with the computed results using a zip that fails on a length mismatch, so a spec that stops describing its graph raises at that point instead of shifting gradients.

Deciding what runs on the device

Once a graph is in the IR, the compiler has to decide, for every node, whether the backend runs it on the device or leaves it on the host, and which of the claimed nodes run together. These decisions set the cost of a compiled call more than any single kernel does, because every edge between a host node and a device node is a transfer across PCIe.

One table for every fact about an operation

The compiler needs several facts about each operation.

• Whether it is elementwise, so it can share a loop with its neighbours.
• How many arguments it takes.
• What it computes, as an expression a code generator can emit.
• Whether the backend has a kernel for it, and for which element types.
• What shape its result has.

Facts kept in separate hand-written tables drift apart, one table listing an operation the other does not, and the symptom is silent. Chains that should fuse do not, and nothing reports it. So every operation is one row in a data file, ops.toml.

[[op]]
name = "add"
aten = ["aten.add.Tensor", "aten.add.Scalar"]
category = "elementwise"
arity = 2
expression = "($0 + $1)"
kernel = true

name is the short spelling hand-written IR uses, and aten lists the spellings PyTorch emits, so a pass can be tested without a model. category says what a pass may do with the operation. arity counts arguments, including constants that arrived as attributes. expression is what the code generator drops into a loop, with $0 and $1 standing for the arguments. kernel says whether the backend runs the operation, and an optional list restricts which element types.

A generator script turns the file into the C++ table the passes read and into the operation coverage page of the documentation. Both outputs are checked in, so a build needs no generator and a reviewer sees what changed, and a test fails if regenerating would change either of them.

Figure 9. Every question the compiler asks about an operation resolves to one row.

Shape rules show the principle at work. Another table could map each operation to a rule, but the rule follows from fields the row already has. A one-argument elementwise operation passes its operand’s type through, a two-argument one broadcasts, a reduction removes or keeps the reduced axes, and a matrix multiply is a matrix multiply. Deriving the rule from the category leaves one table fewer to disagree with the others.

There are six categories, elementwise, reduction, movement, matmul, copy and convolution. Movement covers operations such as view and permute, which rearrange indices and compute nothing. Custom operations add rows at run time to a registry that the same lookup consults, under a seventh category of their own.

Claiming a node

The mark-supported pass records on every node whether the backend will run it, and changes nothing else. A node is claimed when all of the following hold.

• Its operation has a row that says the backend runs it.
• Every value it reads and writes has an element type the row accepts.
• If it is elementwise or a matrix multiply, it has no argument its kernel would ignore.
• Every extent it touches is known at compile time, unless its kernel takes extents as arguments.

The third condition exists because generated kernels read positional arguments only. A keyword argument such as alpha in torch.add(x, y, alpha=2), beta in addmm, or approximate in GELU would be silently dropped, and the kernel would compute a different function. Such a node stays on the host.

The second and fourth conditions reflect one rule. Whether the backend can run something is a property of the node, not of the operation. A matrix multiply is in the table, but cuBLAS is called for float32 only, so a float64 multiply stays on the host. A convolution is claimed only when it is two-dimensional and its extents are known, because its cuDNN plan is built at compile time. A fused loop that has to step through a permuted operand axis by axis needs every extent written into its source, so it cannot be claimed when extents are unknown.

Claiming a node the compiler then cannot build a kernel for is the worst outcome available. The node is placed on the device, so its inputs are copied there. It runs through the CPU fallback, so they are copied back to the host, and its result returns to the device. It pays for every transfer and runs on the host anyway. The defence is to have the claim and the kernel builder ask the same question through the same code. For elementwise operations, whether a loop can read an operand is answered by one function that both call, so they cannot disagree.

Inserting copies where a kernel needs packed data

Generated kernels read their inputs at a flat index, which is correct only for packed row-major data. A claimed elementwise operation or reduction that reads a strided value, such as the result of a permute, needs a packed copy first. The layout-assign pass inserts a node that stands for that copy, which it can do because the value’s layout is in its type.

The inserted node records where the copy happens, so it is visible in the IR and to clustering, which runs next and is described below. It is not claimed, so clustering treats it as a node outside any cluster, which also stops fusion across the copy, as a packed copy would anyway. The executor makes the copy itself, since it packs the inputs of any kernel that expects packed buffers, so the node never needs a kernel of its own. A claimed movement operation needs no copy, since it produces a view, and neither does a matrix multiply, whose library call is told each operand’s layout.

Clustering without cycles

The cluster pass groups neighbouring claimed nodes into units that run on the device together. Grouping matters for two reasons. Transfers are paid at the edges of a unit, since operands cross once when a unit starts and results cross once when it ends. And operations can only fuse into one kernel when they are in the same unit.

The obvious implementation is union-find over edges, merging the two ends of every edge whose ends are both claimed. It takes a few lines, runs in near-linear time, and produces graphs that cannot run. Take four nodes. x adds the two inputs. y reads x and is an operation the backend does not run. z multiplies x by an input. w subtracts z from y. Merging along edges puts x and z together, then z and w, so all three form one unit. That unit contains x, which y reads, so it has to run before y. It also contains w, which reads y, so it has to run after y. No order satisfies both. Each merge joined two claimed neighbours and looked legal on its own. The problem appears only when the unit is treated as a single node of a coarser graph, which can have a cycle even when the original graph did not.

Figure 10. The four-node graph, and two ways to group it. Merging along edges creates a cycle between a unit and an unclaimed node. Keeping only merges that leave the graph of units acyclic gives two clusters that can be scheduled.

So the pass tries each merge and keeps it only if the graph of units is still acyclic. Every node starts in its own unit. Unclaimed nodes are never merged, but they take part in the cycle check, which is what makes the problem visible. A merge that closes a cycle is undone, and the pass repeats until a round merges nothing, because one merge can make another legal. Each round costs time quadratic in the graph’s size, and rounds repeat, against union-find’s near-linear time. Model graphs have hundreds or thousands of nodes, not millions, and what the check prevents is a graph that cannot run at all. A production partitioner would maintain reachability incrementally and reach the same answer faster.

Keeping pure data movement on the host

Movement operations are cheap to claim. On a device tensor they change metadata over the same storage, and mark-supported claims them. Claiming them is not the same as placing them, though. Movement operations tend to sit between operations the backend does not run, whose results are on the host. A unit that contains only a view copies its input to the device, rearranges metadata, and copies the result back, two transfers for no work.

So cluster ends with a rule about units rather than operations. A unit in which every operation is a movement operation is left unplaced, and it runs on the host with the values it already has. A movement operation joins a cluster on the device only when something in the same unit computes. A copy counts as computing, since it is a pass over the data, which is why copy is a category of its own rather than a kind of movement. Filed as movement, every clone would be left on the host by a rule written for views.

The boundary is the cost model

The placement rule is one case of a more general one. Claiming an operation pays when it removes a crossing between host and device, and costs when it adds one. A coverage percentage does not say which. In a ResNet residual block, each convolution is followed by a batch norm. Decomposing batch norm into operations the backend already ran, while the convolutions were still on the host, raised the share of claimed operations and made the block slower, because the boundary moved by one step and nothing stopped crossing. With the convolutions on the device as well, the crossings went away and the block became faster than either earlier version. Placement decisions are therefore judged by the crossings they remove, whatever that does to the number of operations claimed.

The compile pipeline

Every compiled graph goes through the same passes in the same order.

Table 2. The compile pipeline. Each pass documents its precondition, and the order follows from them.

Fusion and code generation

Most operations in a model are elementwise, and fusing them is the first optimisation a graph compiler makes. An elementwise kernel does almost no arithmetic per byte it moves. Run one operation at a time, each writes its result to device memory and the next reads it back. Fused into one kernel, the intermediate values stay in registers and never touch memory.

Why fusion pays

Take four operations on three float32 arrays of 1024 elements, 4 KB each. A multiply of a and b gives m, adding c gives s, a ReLU gives r, and a sigmoid gives the result g.

Table 3. Memory traffic of four separate kernels against one fused kernel.

That is two and a half times less traffic and a quarter of the launches, and the traffic is what sets the time. The fused chain does seven arithmetic operations per element, counting the negate, exponential, add and divide inside the sigmoid, for 16 bytes moved. That is more than a hundred times below the arithmetic intensity at which an RTX 5080 stops being limited by memory, so the kernel’s run time is the time to move its bytes.

Which operations can share a kernel

An elementwise operation computes output position i from its inputs at position i. Two of them in a row read and write the same positions, so one loop can do both. An operation that reads positions it is not writing, such as a sum or a matrix multiply, cannot join the loop.

The fuse-elementwise pass runs inside the clusters and merges a producer into its consumer when all of the following hold.

• Both are elementwise with one result.
• The loop can reach every operand of both. The next section covers operands of a different shape.
• Both are in the same cluster, since a cluster is one unit of execution.
• Both have the same result type, since one loop writes one number of elements.
• The producer’s result has exactly one reader and is not returned by the graph.

The last condition is where fusion stops paying. A value with two readers has to be in memory for at least one of them, so fusing it into the other saves nothing and computes it twice. A value the graph returns is read by the caller from memory, and one that crosses a cluster boundary leaves the unit, so those end a group as well.

Grouping itself is a plain union-find. Clustering needed a cycle check for every merge, and fusion does not, because of the single-reader condition. Every node in a group except the last has its only reader inside the group, so values leave a group only from its last node, and a path from that node back into the group would be a cycle in the original graph, which cannot exist.

One more rule is easy to miss. The natural way to write a fusion pass assigns groups only to chains of two or more operations. If kernels are built per group, an elementwise operation that fused with nothing gets no group, no kernel, and runs through the host fallback. An elementwise node alone in its cluster needs a group of one.

Operands of a different shape

Not every operand of a fused loop has the result’s shape. Layer norm subtracts one mean per row from every element of that row, and a linear layer adds one bias per column to every row. The loop runs one thread per output element, so each operand needs an index expression in terms of the thread’s position i. There are four ways a loop reads an operand.

Table 4. How a fused loop reads an operand.

One function answers “how does a loop over this result read this operand”, and fusion, claiming and the code generator all call it. If each had its own rules, fusion could form a group the code generator then refused, and the group would silently get no kernel. For shapes known at compile time the divisor d is a constant in the source. Shapes that arrive with the call covers what happens when it is not known.

From a group to source code

A group becomes source in two steps, kept deliberately apart. The first resolves the group into a description of the kernel, the values it reads and how it indexes each, the one value it writes, and the list of operations in between. The second, an emitter, turns the description into code for one target. tessiq has one emitter, for CUDA. A second target would be a second emitter receiving the same description, not a second branch inside a growing function, and two emitters stay comparable because they differ only in the text they write. This is the seam where a port to other hardware does its work.

The CUDA the emitter writes is plain on purpose. One thread handles one output element, in blocks of 256 threads, with a bounds check because the grid rarely divides the element count exactly. The element count is a kernel argument even when it is known at compile time, so launching a kernel works the same way whether its shape was fixed or not. Every pointer is marked as not overlapping any other, so the compiler can keep loaded values in registers instead of reloading them after each store. The variables in the kernel carry the names of the IR values they hold, which SSA already makes unique, so a kernel that produces a wrong number can be read against its graph line by line. The command-line driver prints this source for every fusion group, which makes code generation as easy to test as a pass.

Constants go where the IR recorded them. The importer keys a non-tensor argument by its position, and 1 - x, which backward graphs produce often, arrives as a subtraction whose constant is its first argument. An emitter that appended constants after the operands would write x - 1, with the right shape, the right type and the wrong value. Arguments are rebuilt position by position.

Writing math that works for every element type

The element type is a typedef called scalar_t, and math functions are spelled generically, exp rather than expf, so the same expression text in the operation table serves every element type. That has a trap. C++ overloads make exp of a float call the single-precision version, but only if the argument is still a float, and C++ reads a literal with a decimal point and no suffix as a double. In exp(-1.702 * x) the product is a double, and so is everything computed from it. The kernel still returns the right answer and nothing reports the problem, while a consumer GPU runs double-precision arithmetic at a small fraction of its single-precision rate.

The fix is a rule and a test. The code generator writes every constant as scalar_t(...), and a test compiles a generated float kernel to PTX (the assembly-level form the driver consumes) and fails if any double-precision instruction appears in it.

Compiling while the program runs

The source exists as a string inside a process that is already running a model, so it is compiled there with NVRTC, CUDA’s runtime compiler.

NVRTC can produce PTX, which is portable but which the driver compiles to machine code every time it loads it, or a cubin (machine code for one architecture). tessiq asks for a cubin and caches it on disk, under the user’s cache directory. Caching PTX would save the NVRTC compile but still pay the driver’s compile on every load.

The cache key is a 64-bit hash of the source, the target architecture and NVRTC’s own version. The architecture is in it because a cubin loads only on its own architecture or a later one with the same major version, so a cubin built for one generation of card does not load on the next. The compiler asks the driver what the card actually is, and the answer feeds both the compile flag and the key. The NVRTC version is in the key because a toolkit upgrade can change what the same source compiles to, and reusing an old result would make any report about generated code impossible to reproduce.

Figure 11. From a fusion group to a running kernel.

A compilation error returns NVRTC’s own log rather than a summary of it, because nobody wrote the source by hand and the log’s line numbers are the only way into it. A machine without the CUDA toolkit reports that it has none, and a cache directory that cannot be written logs a warning and returns the kernel anyway.

Loading and launching go through CUDA’s driver API, while the allocator and PyTorch use the runtime API. The two cooperate through the device’s primary context. The runtime creates it, and the driver API acts on whichever context is current on the calling thread. Selecting the device through the runtime makes its primary context current, so every loader entry point selects the device first, and a launch from a thread PyTorch has never used still finds a valid context. Modules are loaded once and kept for the life of the process, because loading walks the whole image and a model launches the same kernel many times. Nothing unloads them, since unloading during static destruction at exit races the driver’s own teardown.

How the generated code compares

The generated kernels are simple, and on data that has to come from DRAM that is enough. On a chain of four elementwise operations over 64 million elements, CUPTI (NVIDIA’s profiling interface) times tessiq’s kernel at 1270 µs and Inductor’s at 1277 µs, both at 88% of the card’s DRAM bandwidth. On data that stays in the L2 cache, Inductor’s kernel is 1.3 times faster. The likely reason is that its Triton kernel handles two elements per thread with 8-byte vector loads where tessiq’s handles one element with 4-byte loads. Wider loads per thread are the next step for the emitter, and the separation between description and emitter means it is a change to one class.

Reductions, copies and library calls

Fused elementwise loops cover most of the operations in a model, but not most of its arithmetic. A transformer spends its time in matrix multiplies, a convolutional network in convolutions, both need reductions for their normalisations and softmaxes, and copies appear wherever a transposed tensor has to be laid out again. Each needs a different kind of kernel, and two of them should not be written by the backend at all.

One interface for every kind of kernel

A fused loop, a reduction, a cuBLAS multiply, a cuDNN convolution and a custom operation all tell the executor the same things, so they implement one interface.

• Inputs and output, as IR value names. The importer took them from the FX node names, so the executor can find the tensors a kernel reads and the node it stands in for.
• Members, every value the kernel computes on its way to its output. For a fused loop that is the whole group, and the executor skips the rest of the group’s nodes.
• Whether it packs its inputs. A kernel that assumes packed buffers has its inputs packed by the executor before the launch. A matrix multiply is told each operand’s layout instead.
• Whether its extents are dynamic, and its output shape computed from the tensors it will read, for graphs whose shapes are not fixed.
• Launch, given the tensors, the output and an explicit stream.

The stream is a parameter rather than an assumption so that a kernel written outside the compiler can be told where its work goes. With this interface the executor needs no knowledge of kinds, and adding a new kind of kernel means writing one class, with no new list, binding or branch in the recipe.

Reductions

A reduction combines many elements into one, as a sum, a mean, a maximum or a minimum does. It breaks the rule every fused kernel follows, one thread per output element reading position i of each input, because each output reads many positions.

What lets one kernel cover reductions of every rank is describing each as three numbers. Any reduction over a contiguous run of axes can be viewed as an input of shape outer by reduce by inner, producing outer by inner. outer is the number of independent slices, reduce the number of elements combined into each output, and inner the distance in memory between two elements that are combined. Summing a 4 by 8 matrix over its second axis is 4, 8 and 1. Over its first axis it is 1, 4 and 8. Summing everything in a 4 by 8 by 3 tensor is 1, 96 and 1. The kernel takes the three numbers as arguments and never looks at the rank.

The decision that sets performance is how threads are laid over those numbers so that the threads of a warp read neighbouring addresses and their loads coalesce. Which layout achieves that depends on inner.

Figure 12. The two thread layouts for a reduction, chosen by inner.

When the reduced axis is last, its elements are adjacent in memory, so a whole block of 256 threads takes one output. The threads read adjacent elements together, each keeps a partial result, and the block combines the partials in shared memory, halving the number of active threads at each step. When a kept axis follows the reduced one, it is the outputs that are adjacent, so each thread takes one output and loops over the reduced axis alone, and at every step the threads of a warp read adjacent addresses. Maximum and minimum test for NaN first, so they propagate it as PyTorch does, and a mean divides the sum by reduce at the end. Where PyTorch accumulates in a wider type than the element, as it does for integer sums, or sums in a tree to limit rounding error, a backend has to make the same choices for its results to compare.

From DRAM, both layouts run level with PyTorch on large inputs, at over 90% of the card’s bandwidth. Where it loses to PyTorch, the cause is the limit of each layout. One thread per output needs enough outputs to fill the card, and a first-axis reduction of a matrix with 4096 columns makes only 16 blocks of 256 threads for 84 multiprocessors. One block per output wastes threads when the reduced axis is short, and reducing 64 elements with 256 threads leaves three quarters of every block idle. A production reduction chooses its layout by the number of outputs and the length of the reduced axis, not only by which axis is last.

The copy kernel

clone, and .contiguous(), which reaches it, turn a tensor with any strides into a packed one. Without a device kernel, a copy on a GPU tensor goes through the CPU fallback, the whole tensor to the host and back, for bytes that never needed to leave the card.

The kernel makes two decisions that apply beyond this operation. It moves bytes, not numbers. Rearranging elements never looks at their values, so the kernel is written for an element width rather than an element type. Four bytes are four bytes whether they hold a float or an integer, which covers half-precision, bfloat16 and boolean tensors as well, types the expression emitter has no spelling for. And it takes the shape and the strides as arguments, so one compiled kernel handles every layout of every shape at a given rank and width, including shapes not known until the call. Each thread owns one position in the packed result. It takes its flat index apart into a coordinate per axis, from the last axis to the first, multiplies each coordinate by the source’s stride for that axis, and reads from the sum. Only the source needs that arithmetic, because being packed is what makes the destination’s position equal to the thread’s own index. Kernels exist for ranks up to eight at each of four widths, and a model uses two or three of them. Timed against PyTorch’s own CUDA copy on permuted tensors, it is level at every size.

The same kernel packs the inputs of generated kernels where layout assignment marked a strided operand, since the executor packs them by calling .contiguous().

Matrix multiplies through cuBLAS

A competitive matrix multiply (a GEMM) is a large piece of engineering. It tiles the matrices into blocks that fit in on-chip memory, schedules the tiles across the multiprocessors, and on recent GPUs feeds tensor cores. Vendors spend years on theirs, so a backend’s job here is a binding that describes the problem precisely and lets the library choose the algorithm. mm, addmm (a matrix multiply plus a bias row, and what nn.Linear becomes) and bmm (a batch of multiplies of the same shape) go to cuBLAS in float32.

The description is where the layout in the IR’s types pays off. BLAS takes a transpose flag for each operand, so a matrix stored one way and used the other costs nothing. nn.Linear produces exactly that case on every call, because it stores its weight as out by in and multiplies by the transpose. The binding reads a transpose flag and a leading dimension (the step between rows) for each operand from its strides, so every weight goes to cuBLAS as it is stored. A matrix multiply is the one kernel that does not pack its inputs, since packing a transposed operand would leave the flag describing data that had moved.

BLAS reads matrices in column-major order, and PyTorch stores them row-major. A row-major matrix read as column-major is its transpose, so asking cuBLAS for the transposed product, B transposed times A transposed, with the operands swapped, computes A times B in row-major without copying anything.

For the bias, the binding uses one code path for all three operations. It calls the multiply with a beta of 1, which adds the product to what the output already holds. For addmm a small generated kernel first writes the bias into every row of the output, and for mm and bmm the output is zeroed first. A batch uses cuBLAS’s strided batched call, with each matrix following the last. The binding keeps one cuBLAS handle per device, a handle holds one stream at a time, and handles are shared, so setting a handle’s stream and queuing the multiply happen under one lock. Otherwise two threads could interleave and queue a multiply on each other’s stream.

Convolutions through cuDNN

A convolution can be computed in several ways, and which is fastest depends on the shapes, the filter, the stride, the card and the library version. cuDNN exists to know that, so convolution is a binding too.

cuDNN does not take tensors. It takes descriptors, one each for the input, the output and the filter, giving element type, layout and extents, and one for the convolution itself, giving padding, stride and dilation on each spatial axis, the number of channel groups and the mode. Then it takes raw pointers. None of the descriptors depend on the values being convolved, so tessiq builds them once, when the graph is compiled, into a plan, and each call is a single cuDNN function plus a bias addition when the convolution has a bias. The plan covers two-dimensional convolutions over packed tensors of known shape, and anything else stays on the host.

Two details in the description produce a wrong answer rather than an error when they are wrong. The first is the mode. The mathematical definition of convolution flips the filter, and every deep learning framework, PyTorch included, computes cross-correlation and calls it convolution, so the plan asks cuDNN for cross-correlation. The second is the output shape. PyTorch worked out the output shape while tracing, and cuDNN computes it from the descriptors. They should agree, since both come from the same stride, padding and dilation. If they do not, the backend read an attribute differently from PyTorch, and cuDNN would write a different number of elements into a buffer sized for the graph’s. So the plan asks cuDNN for its answer and refuses to build when the two differ.

cuDNN can choose an algorithm by timing every candidate on the real shapes, or by asking a heuristic that ranks them without running anything. Timing is right for a model deployed once and run for weeks. For a compiler it would run on every compile. tessiq takes the heuristic’s ranking and uses the first algorithm whose workspace (the scratch memory some algorithms need) can be allocated.

A default not to inherit

Recent NVIDIA GPUs can multiply matrices on their tensor cores in TF32, a format with float32’s exponent and range but 10 bits of mantissa instead of 23. It is much faster and agrees with float32 to about three decimal digits. PyTorch lets cuDNN use it for convolutions by default and does not let cuBLAS use it for matrix multiplies. tessiq asks cuDNN for plain float32 arithmetic. Its tests compare results with PyTorch on the CPU, most at a tolerance of 1e-5, and TF32 would fail them while being perfectly acceptable for training a model. Either choice is defensible if it is made on purpose. It also changes every benchmark, since comparing the backend’s convolution with PyTorch’s on the same GPU compares float32 with TF32 unless one side is set explicitly.

Recipes and the executor

The passes decide where each operation runs and which ones share a kernel, and the kernel builders produce the kernels. Something still has to run the graph every time the model is called, and do as little work per call as it can. tessiq splits that job in two. A recipe records everything that depends only on the graph, worked out once. An executor uses the recipe on every call and does only what depends on the tensors.

What is decided once

Building a recipe runs the whole compiler on an FX graph, import, the pass pipeline, kernel construction and memory planning. The recipe holds the following.

Table 5. What a recipe holds.

The move list is the field that changes how a call runs. A node needs its inputs moved only when something it reads is on another device. In the middle of a cluster, every input was produced on the device, so nothing needs checking. Deciding this at compile time means only the nodes on a boundary are inspected during a call, instead of every argument of every node. What is left for each call is launching kernels, dispatching the nodes that have none, allocating results, and copying inputs across a boundary.

A recipe holds names, never pointers into the graph it came from. A cached recipe outlives the graph object it was compiled from, and a pointer into a freed graph reads memory that keeps working until the allocator reuses it. A test builds a recipe, destroys the graph, and reads the recipe back.

A tool prints a recipe for inspection. For tanh(sigmoid(relu(a + w)) @ w), with w a parameter, compiled for a GPU, the kernel section and the schedule read as follows.

== kernels ==

%mm  gemm, extents known here
    reads   sigmoid, arg0_1
    covers  mm
%sigmoid  fused, extents known here
    reads   arg1_1, arg0_1
    covers  add, relu, sigmoid
%tanh  fused, extents known here
    reads   mm
    covers  tanh

== schedule ==

  add         aten.add.Tensor         kernel %sigmoid (moves its inputs)
  relu        aten.relu.default       kernel %sigmoid
  sigmoid     aten.sigmoid.default    kernel %sigmoid
  mm          aten.mm.default         kernel %mm (moves its inputs)
  tanh        aten.tanh.default       kernel %tanh

Five operations become three kernels. The add and the ReLU are computed by the kernel that writes the sigmoid, so they are fused away. The add and the matrix multiply read graph inputs, which arrive on the host, so they are on the move list.

A memory plan

The compiler also plans where every intermediate value could live in one buffer, from each value’s live range, with a greedy first fit that gives each value the lowest offset nothing else alive at the same time occupies. The ranges are closed at both ends. A node may still be reading its inputs while it writes its output, so a value last read by a node cannot share bytes with the value that node writes, unless that particular kernel is known to work in place. Returned values and parameters are left out, since the caller owns them, and the test that matters checks that no two values alive at the same time share a byte. tessiq computes and reports the plan and allocates the real buffers through the caching allocator, which already reuses freed blocks of the same size. The plan becomes the allocation once the executor owns its output buffers, discussed in What comes next.

Compiling a graph once

Building a recipe is not free, and the same graph always produces the same recipe, so recipes are cached. The key has to include everything the compiled result depends on. That means the target device and the rewrites the compile asked for, and, for every FX node, its kind, its target and its arguments, together with the dtype, the shape and the strides of every value a kernel is built from. The key is computed from the FX graph, not from the IR, because importing into the IR is most of what compiling costs, and a key that needed the IR would save almost nothing. One walk over the graph is the least a key like this can cost.

PyTorch caches compiled code too, so why does a backend need a second cache? PyTorch’s is keyed on the Python code object Dynamo traced, so the same graph reached from two call sites compiles twice, and torch._dynamo.reset(), which tests call constantly and some serving code calls between requests, drops it entirely. The backend’s cache is keyed on the graph itself. It holds a bounded number of recipes, 64 in tessiq, and evicts the least recently used, because a recipe keeps its FX graph alive and an unbounded cache would keep every graph a process ever compiled. It is locked, so compiles on several threads can share it.

The executor

The executor is a subclass of FX’s Interpreter, which walks an FX graph node by node and calls a method for each. It overrides three of those methods, run_node, which is called for every node, call_function, which runs one operation, and output, which produces the graph’s result.

Figure 13. What the executor does with each node.

run_node skips the nodes a kernel computes on the way, launches the kernel at the node it stands in for, and otherwise records which node is running, because call_function receives the operation and its arguments but not the node, and placement is a property of the node. call_function moves tensor arguments to the node’s device when the recipe says it must and leaves scalars, dtypes, shapes and None alone. output rebuilds the returned tuple from the output spec.

A kernel’s output is allocated from the shape PyTorch recorded on the node while tracing, or, when that record is symbolic, from the shape the kernel computes from its actual inputs. Inputs are packed first unless the kernel reads layouts itself. Every launch goes to the legacy default stream, which orders it against the copies on each device’s own stream, so launches never wait for one another on the host. Copies do wait. Moving an input to the device is a blocking copy, and so is bringing the result back.

Wiring kernels into the interpreter exposes one assumption. FX’s interpreter frees a value once its last reader has run, and works out the last reader from the graph. A fused kernel reads its inputs at the node that writes its output, later than the last use the interpreter can see, because the uses in between belong to nodes the kernel absorbed. So when a recipe has kernels, the executor keeps values for the whole call. Computing each value’s last use from the kernels’ actual inputs instead of from the graph would free them earlier.

Keeping weights on the device

AOTAutograd lifts a module’s parameters into graph inputs, so every call hands the compiled graph its weights as arguments, on the host. Copied on every call, a 512-wide transformer encoder block would move 12 MB of weights beside 2 MB of input. So the executor keeps the device copy of any input that has not changed since it was copied, in a table that lives as long as the compiled function.

Identity alone is not a safe key. An optimiser writes parameters in place, and a stale device copy would train against old weights while the loss still moved, which looks exactly like a bad learning rate. PyTorch gives every tensor a version counter that increments on every in-place write, so an entry is reused only when the tensor is the same object, at the same storage address, with the same version. The table holds a weak reference to each source tensor, because a freed tensor’s identity can be reused by a new one, and a dead reference misses instead of matching the wrong tensor.

Shapes that arrive with the call

Call a compiled model with a batch of 8 and then a batch of 16, and PyTorch hands the backend a new graph with the numbers taken out. A backend that compiles only for fixed shapes either recompiles for every new size or declines the graph and runs it uncompiled. tessiq compiles the symbolic graph once and runs it at any size.

What Dynamo hands over

By default Dynamo compiles for the concrete shapes it first sees. When a later call changes a dimension, it recompiles, and this time it replaces the dimension that changed with a symbol, on the reasoning that a dimension that changed once will change again. Any backend used on real inputs meets symbolic graphs sooner or later.

A symbolic graph is stranger than a graph with some sizes missing. For a linear layer followed by tanh, called on an input of batch by steps by 32, the batch size and the sequence length arrive as graph inputs, integers PyTorch calls SymInts. There is arithmetic on sizes in the middle of the dataflow, multiplying the two so the three-dimensional input can be viewed as a matrix for the multiply. And a reshape no longer takes a constant list of sizes. It takes values the graph computes.

Representing sizes nobody knows yet

The importer gives every extent PyTorch knows only symbolically the IR’s unknown extent, types the size inputs and the arithmetic on them as rank-0 integers, and turns a size list that contains computed values into operands, with an attribute marking the list as dynamic so no pass reads it as a constant. The size arithmetic is left unclaimed and runs on the host, as a handful of integer operations.

One property of this representation shapes everything after it. Every unknown extent is the same question mark. The IR does not record that two unknown extents are the same symbol, so it cannot prove they are equal. Kernels therefore check at launch, with real sizes, what the compiler could not check.

Support depends on the shape

Whether the backend claims a node now depends on whether its kernel can take its numbers at launch.

Table 6. What is claimed when a node’s extents are not known until the call.

The rule behind the table is the one from Claiming a node. A node whose kernel needs known extents, claimed anyway, gets no kernel and runs through the CPU fallback with its inputs on the device, which is worse than leaving it on the host.

Where a kernel gets its numbers

A kernel compiled for a symbolic graph runs at sizes nobody knew when it was compiled, so every number it needs arrives with the call. The executor has the tensors, and each kind of kernel derives what it needs from them.

Fused loops already take their element count as an argument, read from the output tensor. A broadcast operand needs one more number, the divisor d, which becomes a kernel parameter when it is not known at compile time. At launch the executor has pointers and element counts and nothing else, and that is enough, because a broadcast is a relation between element counts rather than a fact about shapes. A repeated operand, read at i / d, holds one element for every d elements of the result, so d is the result’s element count divided by the operand’s. A cycled operand, read at i % d, holds exactly one cycle, so d is its own element count.

Both divisors are also checkable, and the checks stand in for the equalities the IR cannot prove. Before a launch, every operand read at i must have exactly the result’s element count, and every broadcast operand’s count must divide it evenly, or the last run or cycle would read past the end of the buffer. A buffer that fails is refused with an error instead of read out of bounds. Where a divisor is known at compile time, as a layer norm’s row length is when the model’s width is fixed, it is still written into the source as a constant, because dividing by a constant is cheaper than dividing by a value in a register.

Reductions already take their three extents as arguments. The compiled description records which axes are reduced and the input’s rank, and the launch computes the extents from the actual tensor.

Matrix multiplies need nothing new from cuBLAS, which has always taken its dimensions and leading dimensions as call arguments. They now come from the tensors at the call.

Output sizes are the last number. The executor allocates each kernel’s output before launching it, normally from the shape PyTorch recorded while tracing. In a symbolic graph that record is symbolic, so the executor asks the kernel, which computes the result’s shape from the tensors it will read. A fused loop’s output is its inputs’ shapes broadcast together, a reduction’s is its input with the reduced axes removed or kept, and a matrix multiply’s is the rows of one operand and the columns of the other.

What it costs

One symbolic compile then serves every size. Twenty-one different sequence lengths through a 512-wide encoder block caused two backend compiles, one for the first concrete shape and one symbolic compile when Dynamo saw the length change, and every later length ran the same recipe. The symbolic recipe gets the same kernels as a static one. A call costs about 17% more than a call to a compile made for that exact shape. The split was not profiled. The candidates are the size arithmetic on the host, the kernels computing their output shapes, the launch-time checks, and division by values in registers where a static kernel divides by constants.

Custom operations

A compiler generates good code for the patterns it knows. Sooner or later a model needs an operation written by hand, either because the compiler cannot generate it or because a hand-written kernel does better. RMSNorm, a fused softmax and a GEMM with its bias and activation folded in are typical. PyTorch lets anyone define a new operation, but on its own that gives a backend nothing. The compiler has never heard of the operation, cannot claim it, and runs it on the host between two device kernels. A backend needs a way to add operations that its compiler, its eager path and its tests all understand.

An operation is a directory

In tessiq each custom operation lives in its own directory in the C++ source tree. The definition is a C++ file. It gives the operation’s schema, its signature in PyTorch’s own schema syntax. It gives a reference implementation written with ATen (PyTorch’s C++ tensor library). It gives a launch function, which receives the call’s arguments and the output and returns a launch configuration, meaning the grid and block dimensions, the dynamic shared memory, and the kernel’s scalar arguments. And it lists the element types the kernel handles. The scalar arguments are typed as the C types of the parameters they fill, because a kernel reads its arguments as raw bytes in the order its signature declares, and a double handed to a float parameter is read as the wrong number rather than refused. The kernel itself sits beside the definition in a .cu file, with one pointer parameter per tensor argument, then the output, then the launch arguments.

Adding an operation is a new directory and a rebuild. The build finds every directory, compiles each definition into the extension, and embeds the text of each kernel file in it. The kernel is compiled with NVRTC for the GPU it runs on, the first time it runs there, and cached on disk like the compiler’s own kernels. One build therefore works on every generation of card, and a machine with no GPU runs the reference. Keeping definitions in C++, beside the backend’s other operations, means a compiler checks them before anything runs, which a definition written from Python would not get.

Three kinds of kernel

Table 7. The three ways a custom operation provides its implementation. Each operation declares exactly one.

An operation that computes each output element from the inputs at the same position needs no kernel at all. It gives an expression, in the same form as the operation table’s, and the compiler fuses it into the loops it generates. The table’s lookup checks the generated table first and then a registry of operations added at run time, so claiming, clustering, fusion and the code generator treat a registered expression exactly as they treat a built-in one. A custom fast GELU between a multiply and a ReLU compiles to one kernel with both. Registration refuses a floating-point constant in an expression written without scalar_t(...), for the reason given in Writing math that works for every element type.

Checked when the package is imported

The definition macro records an operation when the extension loads and does nothing else. The checks run when the package is imported, because a C++ exception thrown while a shared library is loading ends the process, Python included, rather than raising anything Python can catch. At import, a bad definition fails the import with a message that names the operation and the problem. The checks cover a schema that does not parse, a reference or launch function whose parameters do not match the schema, a missing kernel file, an operation that declares more or fewer than one kind of kernel, a bare floating-point constant in an expression, and the replacement and pattern rules described below. A separate check function compiles every kernel, which the test suite runs for each operation.

What the reference is for

The reference is registered for every kind of tensor and does three jobs.

It runs wherever the kernel cannot, on the host, on the CPU reference device, and for an element type the kernel was not written for. It is the oracle every kernel test compares against. And it tells the backend the shape of the result. A meta tensor has a shape and a dtype and no data, and running an ATen operation on meta tensors computes the result’s shape and dtype without touching memory. So the backend runs the reference on meta tensors to size the output, and tracing under torch.compile does the same. An operation author never writes a shape function, and a shape that changes from call to call works because the reference’s own arithmetic on sizes is already symbolic.

The reference also creates a testing problem. A call on the GPU gives the same answer whether the kernel ran or the call fell back to the reference on the host, so a test that only compares results passes either way. Each operation therefore counts its fallbacks, and a test asserts the count. A float32 call on the GPU must leave it unchanged, and a float64 call must add one. A launch function can also decline a call it cannot handle, such as a weight of the wrong shape, and the call then runs the reference and is counted.

Streams are explicit

Every kernel the compiler generates, and every library call it makes, goes to the legacy default stream, which orders itself against the runtime’s copies. A kernel from outside the compiler that queued its work on some other stream could run before its inputs had been computed, and would be wrong only some of the time. So the stream is an argument all the way down. Every kernel’s launch takes one, and so does every call that tells cuBLAS or cuDNN which stream to use.

Standing in for a PyTorch operation

A custom operation only helps a model that calls it, and most models call PyTorch. So an operation can declare that it replaces a PyTorch operation, which makes existing models use it unchanged. tessiq’s softmax kernel replaces _softmax, which torch.softmax, F.softmax and nn.Softmax all reach. The schema has to match the PyTorch operation’s argument names and types, and the replacement takes effect in three places.

• Eager calls. The custom operation is registered as _softmax’s implementation for the backend’s dispatch key, so torch.softmax on a GPU tensor of the backend runs the kernel.
• Compiled graphs. A pass named after the operation runs before claiming and turns each _softmax node into the custom operation, wherever its types are ones the operation declares.
• Decompositions. The decomposition table would otherwise break softmax into five operations before any graph reached the compiler. Where a custom operation replaces a PyTorch one, its decomposition returns NotImplemented for the declared types, which tells AOTAutograd to keep the operation whole. For other types it decomposes as before.

Not every PyTorch operation can be replaced. A CompositeImplicitAutograd operation is broken into others as it is called, before any backend kernel or graph sees it. softmax is one, which is why the operation replaced is _softmax, the one it breaks into. Declaring a replacement for a composite operation fails the import with a message saying to replace one of the operations it breaks into. A replacement keeps PyTorch’s backward for the operation, because only the forward implementation changed.

Finding the computation a model writes out

Most models write RMSNorm out in full, as x times the reciprocal square root of the mean of x squared plus epsilon, times a weight. After decomposition that is six nodes, which the compiler turns into four kernels. An operation can declare a pattern, which registers a rewrite pass that finds those six nodes and replaces them with one call. The pattern is a graph in the compiler’s own text format, stored beside the kernel, which makes it as easy to read and to test as any other IR.

The matcher follows a few rules.

• Each pattern input binds to one value of the model’s graph and must bind to the same value every time it appears. A graph that squares one tensor and multiplies another is not RMSNorm.
• Every attribute must be equal, compared numerically, and a pattern can capture a number instead of fixing it. RMSNorm’s pattern captures epsilon, which becomes the operation’s argument.
• Types are not compared by the matcher, and a match whose types the operation does not declare is left alone.
• The two operands of an addition or a multiplication match in either order, since a model is as likely to write one order as the other.
• The pattern’s inputs are named after the operation’s tensor arguments, and the import checks that the pattern binds or captures every argument.

Whether a match may be replaced depends on one more rule. Every value inside the match other than its result stops existing, so none of them may be read outside the match or returned by the graph. That rule is also what keeps training correct without any code for it. A training step’s forward graph returns every value its backward graph will need, the normalised values inside RMSNorm among them, so in a training graph the pattern finds nothing it may replace and the step runs as written.

The replacement keeps the name of the value the pattern returned and records the names of the nodes it absorbed, because the executor walks the original FX graph, not the IR, and has to know which nodes the new kernel computes. A rewrite changes which operations run, and so can change results in the last bits, so it runs only when a compile asks for it by name. On a small block made of an RMSNorm, a linear layer and an approximate GELU written in plain PyTorch, the RMSNorm and GELU rewrites take the block from six kernels to three. At small sizes most of the saving is on the host, as fewer launches and fewer nodes for the executor to walk.

Precompiled operations

Custom operations with a launch function are single CUDA kernels compiled at run time with NVRTC. That covers a kernel and nothing around it. NVRTC compiles device code only. It cannot compile host code that launches several kernels, allocates scratch memory between them, or calls a library, and it cannot compile headers that need a host compiler, which rules out most of CUTLASS and all of cuBLASLt’s API. Precompiled operations close that gap with host code compiled by nvcc when the backend is built.

Declaring host code

An operation declares itself precompiled and puts its host code beside the definition, in a file compiled by nvcc that defines one host function and registers it with a macro. The build compiles every such file for the architectures in CMAKE_CUDA_ARCHITECTURES, which defaults to the GPU in the build machine. A file that includes a CUTLASS header gets CUTLASS, which is header-only and fetched when the build is configured, and a file that includes the cuBLAS or cuBLASLt headers is linked against those libraries. Nothing else in the build pays for either. On a machine without nvcc these files are skipped, and their operations always run their reference.

The call interface

The host function receives a call object, defined in a header with no PyTorch dependency so that host code compiles with nvcc alone.

Table 8. The interface precompiled host code is given.

The output is allocated before the function runs, with the shape the reference gives on meta tensors. Workspace comes from the same caching allocator as every tensor, so scratch memory is pooled and counted, and its lifetime follows the stream rules from the runtime. When the call’s work went to the legacy default stream, the workspace goes back to the allocator as soon as the work is queued, which is safe because any later reuse is queued after it. When a caller handed the operation another stream, which lacks that ordering, the stream is synchronised first. And when that stream is being captured into a CUDA graph (a recording of queued work that can be replayed later), the workspace is kept for the life of the process, because every replay reads it again. A launch error left pending when the function returns fails the call with CUDA’s message, so host code does not have to check every launch.

Refusing the wrong GPU

Host code built by nvcc contains machine code for the architectures it was built for, plus PTX for any architecture it was built for as a virtual target, which the driver can compile for newer GPUs. On a GPU with neither, a launch fails. To report that before anything runs, the registration macro also compiles an empty probe kernel into the file. Whether the CUDA runtime can find an image of that kernel for a device is exactly whether the file was built for that device’s architecture. The check runs before an operation first runs on a GPU, when the test suite checks every operation, and when a compile builds its kernels, and its message names the architecture to add to the build.

Two ways to fold a bias and GELU into a GEMM

A linear layer followed by GELU, compiled normally, is three launches, a bias fill, the cuBLAS GEMM and a generated GELU kernel. Folding the bias and the GELU into the GEMM’s epilogue (the stage that turns each block’s accumulators into output elements) saves writing the product out and reading it back. tessiq has two implementations, one at each end of what precompiled operations allow.

The first is a CUTLASS GEMM with the exact GELU in its epilogue, and its host code chooses everything. It uses plain float32 arithmetic on the ordinary cores, reads the weight as stored with no transpose, and passes the bias as an operand with a zero row stride so that every row of the output reads the same bias. It uses one tile shape, and splits the inner dimension into slices that run side by side when the output has too few tiles to fill the card.

The second calls cuBLASLt with its bias-and-GELU epilogue and lets the library choose. Its host code caches the heuristic’s choice of algorithm per device, shape and pointer alignment, because asking takes longer than running a small GEMM. It also tells the heuristic the real alignment of the pointers, since cuBLASLt assumes 256-byte alignment unless told otherwise and a view at an offset has less. CUPTI shows that for these float32 problems the chosen algorithms run the epilogue as a second kernel after the GEMM, so one library call is not always one kernel.

Measured on kernel execution, the cuBLASLt version beat the unfused baseline at all four shapes tried, by 2 to 14%. The CUTLASS kernel won at one and lost at three, because a single tile configuration only competes where it suits the shape as well as cuBLAS’s per-shape choice does.

Measuring a backend

A backend’s performance numbers are easy to get wrong in ways that still look plausible. A GPU runs asynchronously, so a timer around a launch measures the launch, not the work. A benchmark that reuses the same buffers can run out of a cache rather than memory. And a profiler is a program too, whose cost lands inside what it measures.

A profiler that adds up

tessiq’s profiler records what happens during compiled calls on the current thread, in three kinds of event. A kernel event is the launch of one of the recipe’s kernels, an interpreted event is a node with no kernel dispatched through PyTorch, and a transfer event is a copy between host and device. Each records host time, and a profile can be exported as a Chrome trace (a JSON format that trace viewers display as a timeline).

Events record exclusive time, the time a region spent outside any region measured inside it. A node that moves its inputs and then runs contains a transfer, and with inclusive times the transfer would be counted twice, once as itself and once inside the node. With exclusive times the events add up to the call, so a profile can be read as a budget.

Two clocks for two questions

A kernel launch returns as soon as the work is queued, so a host timer around it measures the cost of asking, a few microseconds whatever was asked for. The obvious fix is to synchronise the device before stopping the timer. That measures the right thing and changes what is being measured. Normally the host queues work while the device works through the queue, and the two overlap. Synchronising after every kernel turns that overlap into a strictly serial sequence, and the reading is not even the kernel’s own time, since a synchronisation waits for everything queued, so each region absorbs whatever was still running from before it began.

So the profiler uses two clocks. The host clock measures what issuing work costs. The device clock comes from CUDA events, markers recorded into a stream that the device timestamps when it reaches them. Two events around a kernel give its device time, and recording an event does not wait. Reading the times does wait, so the profiler reads them once, when the profile ends and the work has finished anyway.

Events have a cost of their own. On the test machine, starting one of the profiler’s device timers costs about 1.7 µs of host time, and launching a small generated kernel about 3.1 µs, so a start and a stop around every kernel adds host work comparable to the launch being observed. Device timing is therefore off by default and turned on when the question is which kernel is slow. Host time answers where a call’s time goes. A measuring tool’s cost lands inside the thing it measures, and the only check is to run with the tool and without it and compare.

Two kinds of device time

Event pairs around each kernel in a real call carry an event’s worth of bias each, which is fine for ranking kernels and wrong for adding them up. A second tool answers the additive question differently. It launches each of a recipe’s kernels many times between one pair of events, on inputs shaped like the real ones, and reports the mean, which spreads the cost of the two events over all the launches.

What that measures is the stream’s time for back-to-back launches, and it is not the same as the kernels’ own execution time. CUPTI records each kernel’s start and end on the device, and torch.profiler exposes those records. With a PrivateUse1 backend imported they are labelled with the PrivateUse1 device type, so a filter on the CUDA type finds nothing. On the test machine the two measures differ by a factor of three for the same kernels, because the device leaves a gap of several microseconds between one kernel’s end and the next one’s start even when the next is already queued. Kernel execution says what the generated code costs. Back-to-back time says the least a call could take if the host issued everything instantly. A number labelled “device time” should say which one it is.

Where every operation went

tessiq.explain prints how a model compiles, in sections. There is a summary, the imported graph, what each pass changed, the kernels, and a schedule that says what happens to every operation. Each operation is marked as launched as or computed inside a named kernel, as interpreted on the device, meaning it runs through the dispatcher on the device with no kernel of its own, or as run on the host, and a node that has to copy its inputs first says so. Those three fates look the same in every other view of a compiled model. The report reads the pass pipeline from the compiler, the kernels from the recipe and the placement from the map the executor uses, and keeps no list of its own.

Benchmarks that measure what they claim

The numbers in a backend’s documentation follow a few rules, each of which exists because the opposite produces a plausible wrong number.

1. Make the data larger than every cache, or clearly smaller. A benchmark that runs the same kernel on the same buffers repeatedly, with a working set that fits in L2, reads mostly from the cache and can report bandwidth above the card’s DRAM peak. A working set close to the size of the cache is worse than either, because how much of it stays cached depends on where the buffers happen to land in memory, and the same kernel can take twice as long from one run to the next.
2. Time both sides the same way. A raw kernel launch timed against a competitor’s whole compiled call, which includes guards and a Python wrapper, compares a kernel with a kernel plus overhead. At small sizes the overhead is larger than either kernel and the comparison says the opposite of the truth.
3. Know what a small measurement is made of. Timing a kernel of tens of microseconds with a host clock and a synchronise measures launch and synchronisation latency as much as the kernel, and varies by half from one run to the next. Use CUPTI’s records of the kernel, or a size at which the kernel dominates.
4. Measure on an idle GPU, and record the load before and after.
5. Measure the baseline in the same session, alternating. Host times on one machine can move by a quarter between sessions, so a before and an after taken on different days compare the days.
6. Say what was timed and how it was summarised. A whole call from host tensors, a kernel with its data on the device, back-to-back device time or kernel execution, and medians or means.

Testing a backend

Most bugs in a numerical backend do not crash. A constant on the wrong side of a subtraction, a gradient on the wrong parameter, a view that lost its aliasing, a kernel that silently computes in double precision, each produces an answer of the right shape and dtype that is wrong or slow. A test suite for a backend is designed around catching that kind of failure, and around proving that the thing it tests actually happened.

A reference to compare against

Compiled results are compared against PyTorch’s own, run eagerly on the CPU, most of them at a tolerance of 1e-5. The same test runs twice, once compiled for the CPU reference device and once for a GPU. On the reference device every claimed operation goes through the CPU fallback, and the test still exercises the importer, the passes, the recipe, the executor, the placement decisions and every transfer, on a machine with no GPU. The tolerance is part of the design, and it is why the backend asks cuDNN for float32 arithmetic, as described in A default not to inherit.

What CI runs

CI runs on machines without a GPU. It builds the C++ core and runs its tests, including the golden files for every pass, the verifier after every pass in debug builds, and the test that runs every pass twice. It runs the C++ tests again under the sanitizers, runs every Python test not marked as needing a GPU, and enforces the formatters, linters and type checker. That covers the compiler’s decisions and the whole compiled path on the reference device. It does not cover the generated kernels as compiled and launched, the library bindings, the copy kernel, precompiled operations, or the GPU versions of the model and training tests. Those carry a marker and run on a machine with a GPU, and a project that can afford a GPU runner in CI should add one.

Proving the work ran where it should

A test that compares results passes whether or not the backend did what the test is about. A compiled model that quietly ran everything on the host gives the same answer as one that ran on the device. So several tests check side effects as well as results. They require the device allocator’s counters to move during a call, assert a custom operation’s fallback count, compile a generated kernel to PTX and look for double-precision instructions, and read the recipe to check which nodes are claimed, clustered, fused away or left on the host.

Testing training

The backward graph is traced with the forward one but compiled only when something first calls .backward(), so a suite of inference tests covers half of the compile path and looks complete. Training tests run optimiser steps through the compiled model and compare against eager. Checking that the loss goes down is not enough, because a backward pass with a small error still trains, to a different set of weights. The tests compare every parameter after two hundred steps, at a relative tolerance of 1e-4, looser than for single results because rounding differences accumulate over the steps.

Backward graphs contain shapes forward graphs rarely produce, and training tests are the only tests that reach them. They return None in the positions of frozen parameters. They contain 1 - x, with the constant first. And a module that computes its own loss starts its backward graph by broadcasting one gradient value to every element, which produces a tensor whose strides are all zero. Each of these needs a test module written to produce it.

A test has to be seen to fail

Every test that proves something happened should be checked by breaking the code on purpose and watching it fail, since a test that passes on broken code hides the bug it was written to catch. The design of the test often has to change to make that possible. A test that a kernel honours the stream it is given captures the launch into a CUDA graph and replays it. A launch that ignored its stream would have run immediately, during the capture, and written the right answer by itself, so the test fills the output with NaN between the capture and the replay. A recorded launch overwrites the NaN on replay, and a launch that ran early leaves NaN behind.

Inputs chosen to break things

Tests written alongside a feature use the inputs the feature was written for. The inputs that find bugs in a backend are the others, and they are worth listing for any new operation or pass.

• Mixed element types in one expression, such as a boolean mask multiplied by a float tensor.
• Integer inputs to reductions, where PyTorch widens the result type.
• NaN and infinities, and negative zero.
• A layout that changes between calls, such as a transposed input where a packed one was compiled.
• Views as destinations, such as copying into a slice of a larger tensor.
• Every dtype the operation table claims, not only float32.

What comes next

A backend built in this order reaches a point where the compiler, the kernels and the tests are in place, and the next stages are clear.

Eager operations on the device. The operations in the table run on the device only inside compiled graphs. Called eagerly, anything other than allocation, copies, views, clone and custom operations still goes through the CPU fallback. The next stage registers eager implementations for the table’s elementwise operations, generated from the same table, each running a one-node recipe built once per operation, dtype and rank and then cached. Elementwise custom operations already run eagerly this way. The stage after it adds eager reductions, matrix multiplies through cuBLAS and convolution through cuDNN.

A faster executor. On small graphs the executor, not the kernels, sets the cost of a call. On the test suite’s 32-wide encoder block, tessiq’s kernels execute for about 12% of a call, and even queued back to back they would take only about 39% of it. The rest is the host walking the graph in Python and dispatching views, permutes and expands at several microseconds each. There are three known options. The loop can move to C++. The views and permutes, which only compute metadata from shapes the compiler already knows, can be resolved at compile time instead of dispatched on every call. Or the whole call’s device work can be captured as a CUDA graph and replayed, which removes the host’s per-launch cost for graphs whose shapes and pointers are fixed. Each needs the recipe to describe a call more completely than it does now, and the memory plan becomes the allocation once the executor owns its buffers.

Reduced precision. Models trained or served in float16 or bfloat16, usually through autocast, reach the compiler as graphs with explicit casts and half-precision operations. The IR carries both types and the copy kernel moves them, but the expression emitter has no spelling for them and the library bindings call cuBLAS and cuDNN in float32, so tessiq does not compile them yet. Until a backend can build a half-precision kernel, it should leave those nodes unclaimed, so that they run on the host without a transfer to the device and back. Supporting them means emitting half-precision loads with float32 arithmetic inside fused loops, calling the libraries’ mixed-precision entry points, and deciding deliberately whether tensor cores may be used, as the TF32 question already requires.

Overlap. Ordering everything through the legacy default stream keeps the allocator and the executor simple, and means transfers never overlap with compute. Pipelining them needs separate streams, events between them, and an allocator that tracks which stream each block was used on.

Custom operations with a backward. A custom operation that a model calls directly needs a backward of its own before the model can train through it, and schemas with default values need the compiled call to fill in arguments the graph dropped.

Porting to other hardware

A port keeps most of the backend and replaces a known list of pieces.

• The device implementation, written against any runtime with an allocator, a memory copy and a way to wait, together with the device kind and its discovery. The caching allocator and the copy operators come with it unchanged.
• Pinned host memory, the device-side copy kernel, and the Python helpers that name a device’s backend, which assume CUDA today.
• An emitter for the device’s kernel language, fed the same kernel descriptions the CUDA emitter receives, and a way to compile and load what it writes in place of NVRTC and CUDA’s driver API.
• Library bindings for the device’s own matrix multiply and convolution libraries in place of cuBLAS and cuDNN, and new builds of any precompiled operations.
• The stream model. The runtime relies on the legacy default stream ordering itself against the copy stream. A device without that implicit ordering needs events between its queues, and an allocator that tracks which queue each block was used on.

The IR, the passes, partitioning, the recipe, the executor, the reference device and the tests carry over, and so does most of the operation table. Its kernel column changes with what the new device can run, and its expressions are written in CUDA C++, so they have to be translated for another kernel language. A device that can run Triton should weigh plugging into Inductor instead of writing an emitter, since that replaces the third item and most of the compiler with work Inductor already does.

The order I built tessiq in is the order to follow on new hardware. Build the runtime and the reference device first, get every operator working slowly through the fallback, then let the compiler take operations over one at a time, with a test against PyTorch for each stage and a measurement before each optimisation.

---

Thanks for reading this far. It is a long article about a long project, and I hope it saves you some of the time I spent working out what goes where. Backends, and software written close to the hardware in general, are the work I enjoy most, which made this a pleasure to write (with help from my fellow AI agents, of course).

tessiq is not finished, as the list of next stages shows, and parts of this article will change as it grows. If you are building a backend for your own hardware, or you find something here that is wrong, I would like to hear about it (LinkedIn).