Introducing TrainKernel

In this article

TrainKernel is an experimental compiler that turns a supported PyTorch training step into one generated CUDA kernel. On tiny online workloads, that kernel can beat PyTorch Inductor. On larger matrix multiplications, it loses badly. Both halves of that result explain what a training megakernel is for.

A normal training step looks simple in Python. Run a forward pass. Measure a loss. Differentiate it. Update the parameters. On a GPU, that expression becomes a sequence of kernel launches, temporary tensors, library calls, and synchronization points. For a large model, the useful arithmetic dominates. For a tiny model or batch of one, the machinery around the arithmetic can become a meaningful fraction of the step.

PyTorch already attacks this problem. torch.compile captures Python programs and uses TorchInductor as its default optimizing backend. Its low-overhead modes can also use CUDA Graphs, which remove repeated host launch work.1PyTorch, torch.compile documentation and compiler overview. TrainKernel compares against Inductor with CUDA Graph capture rather than an uncompiled Python loop.

TrainKernel asks a narrower question: if the complete training step is small, static, functional, and repeated many times, can the forward pass, loss, backward pass, and optimizer update share one persistent device program? The point is not to replace PyTorch. It is to find the boundary where eliminating the remaining device-side boundaries is worth giving up a general compiler's mature kernels.

When the model is small enough, launching the work can matter almost as much as doing the work.

One training step, one kernel#

The public interface is deliberately ordinary. A training function receives detached state and data, computes its own gradients with torch.func.grad_and_value, and returns the loss, gradients, and updated state. TrainKernel captures that function once for fixed shapes, strides, devices, and dtypes.

import torch
import trainkernel

def train_step(weights, x, target):
    def loss(w):
        prediction = torch.relu(x @ w)
        return ((prediction - target) ** 2).mean()

    gradients, value = torch.func.grad_and_value(loss)(weights)
    return value, gradients, weights - 1e-3 * gradients

compiled = trainkernel.compile_training(train_step, (weights, x, target))
loss, gradients, weights = compiled(weights, x, target)

The generated program is not a long list of CUDA launches. It is one cooperative kernel whose resident blocks execute the graph in topological order. Grid-wide barriers separate stages that have real dependencies. Pointwise operations can be inlined into their consumers, temporary buffers can be reused after their final read, and safe last-use parameter updates can commit directly instead of passing through another staging buffer.

The grid-wide synchronization is what lets later parts of the same kernel consume values produced by earlier parts. CUDA only permits this kind of whole-grid barrier under a cooperative launch, where the runtime can guarantee that the participating blocks are resident together.2NVIDIA, CUDA Cooperative Groups. Whole-grid synchronization requires cudaLaunchCooperativeKernel; TrainKernel also checks device support and clamps its grid to measured occupancy.

This architecture creates a trade. It is excellent at removing boundaries. It is much less free to assign every matrix multiplication a separate launch geometry, deep pipeline, or library implementation. The same resident grid that makes a whole step possible also limits how much parallel work a large matrix multiply can expose.

Inside the compiler#

Capture begins with functionalization and fake-tensor tracing. Instead of executing the example update, TrainKernel obtains an ATen graph and a structured description of the returned loss, gradients, and new state. Unsupported mutation, tensor constants hidden in closures, cross-device operations, unsupported dtypes, and dynamic shapes fail at the boundary. There is no silent eager fallback.

The lowered intermediate representation is intentionally small. Every tensor value carries its shape, dtype, root storage, offset, and strides. Views remain aliases when possible. Materialized values receive workspace storage, and a liveness pass lets later values reuse regions whose final reader has executed. The compiler also separates direct updates from transactional updates. A parameter can be overwritten in place only after TrainKernel proves that no later operation needs the old value.

Code generation has three main arithmetic paths. FP32 matrix products use a tiled SIMT implementation to preserve the expected accuracy. BF16 and FP16 products use WMMA tensor-core fragments. Row-shaped products can take a split-K GEMV path. Around them, the compiler emits reductions, broadcast arithmetic, ReLU and its backward, scalar powers, casts, fills, and supported storage views. The complete implementation and design notes are public.3Arya Somu, TrainKernel source repository and compiler design, version 0.2.0 (2026), Apache 2.0.

Benchmarks: complete steps, not isolated operators#

The benchmark is a two-layer ReLU MLP trained with FP32 master weights, BF16 forward and backward matrix compute, FP32 loss, and FP32 SGD updates. The four shapes move from a batch-four tiny MLP to a batch-128 medium MLP. Every timed method performs the complete forward, loss, backward, and update step.

Each method is captured as a stateful 20-step CUDA Graph over four fixed GPU-resident batches. Compilation, capture, reset, warmup, and setup are excluded from the event timing. Every seed retains 30 paired samples, a 100-step numerical trajectory, a 20-step replay check, generated-source hashes, and baseline failures. The reported number is the median of three held-out seed medians. Inductor was allowed multiple declared configurations, and the table compares TrainKernel with the fastest measured one. This is intentionally a hard baseline, not eager PyTorch dressed up as a compiler comparison.4TrainKernel, A100 benchmark methodology, H100 methodology, and the raw held-out JSON files linked from each result report.

Relative latency across held-out workloads

TrainKernel latency divided by the fastest measured Inductor latency. The scale is logarithmic; the vertical line at 1× is parity. Points to the left win.

0.5×1×2×4×8×

A100 SXM4 80 GB

tiny

0.67×

online

0.84×

small

2.89×

medium

6.05×

H100 SXM 80 GB

tiny

0.82×

online

0.62×

small

2.86×

medium

6.26×

An independent one-hot least-squares classification head gave the same directional result. On A100, TrainKernel measured 13.07 microseconds against 18.63 for the fastest measured Inductor configuration, a 1.43× win. On H100 it measured 12.77 against 15.65, a 1.22× win. The workload is useful because it resembles a small online head rather than another seed of the same MLP shape.

The cost appears before the first fast step. The independent A100 classification run took 53.62 seconds to compile and capture from a cold cache. At the measured steady-state latency difference, the report estimated that about 7.82 million updates would be required to amortize that setup. TrainKernel is therefore plausible for a long-lived repeated loop, not for a command that runs for five seconds and exits.5TrainKernel, independent A100 release validation. The report preserves three seeds, correctness failures, raw samples, source hashes, and the setup-amortization calculation.

The H100 experiment: a correct kernel can still be the wrong kernel#

The first H100 result used the same WMMA-oriented backend selected on A100. A launch calibration improved the medium workload from 209.03 to 191.09 microseconds, about 8.6%, but it did not change the larger conclusion. I then added an opt-in SM90a path using Hopper's asynchronous warpgroup matrix multiply instructions, or WGMMA.

WGMMA is an architectural temptation. A warpgroup can issue asynchronous matrix multiply-accumulate work and later commit and wait on groups of operations. In an isolated 64×16×64 probe, a register-source variant was about 6% faster than the shared-source version. That looked promising. In the complete training step, it lost.6NVIDIA, PTX asynchronous warpgroup matrix instructions; TrainKernel, H100 calibration evidence.

H100 calibration: WMMA remained faster

Microseconds per complete step on the fixed calibration seed. Shorter is better. These calibration measurements are separate from the held-out table above.

Small workload

WMMA

56.53 μs

WGMMA g1

93.59 μs

WGMMA g4

92.87 μs

WGMMA reg A

104.85 μs

Medium workload

WMMA

190.88 μs

WGMMA g1

344.73 μs

WGMMA g4

355.29 μs

WGMMA reg A

389.76 μs

This was the most useful negative result in the project. New instructions do not erase surrounding work. TrainKernel still had to stage operands, synchronize the warpgroup, move accumulators, execute pointwise work, and preserve the cooperative schedule. A locally faster matrix primitive made the complete program slower.

A serious Hopper-native successor would need to treat WGMMA as one part of a different pipeline: tensor memory accelerator transfers, swizzled shared-memory layouts, warp specialization, deeper overlap, and perhaps thread-block clusters. Dropping one new instruction into an A100-shaped schedule was enough to prove correctness, not performance. The public release keeps the slower experiments instead of selecting a favorable microbenchmark.

Where a training megakernel could be useful#

The benchmarks support a narrow set of candidate use cases. They do not prove production gains in any of them, but they identify the shape of a workload worth testing next.

• Online learning with small models. A service that performs one or a few updates after each event may care about step latency more than bulk throughput. Static feature widths and long-lived state fit TrainKernel's current contract.
• Continuously adapted classification heads. The independent least-squares head is the closest implemented example: a small resident head, repeated updates, and no data-loading time inside the measured loop.
• Bandit and control loops. A small policy or value model updated inside a simulator can run millions of structurally identical steps. The setup cost can be amortized, and a few microseconds may matter to the outer loop.
• Inner-loop adaptation experiments. Fixed-shape meta-learning or per-task adaptation can contain tiny repeated training programs. The current operator subset is restrictive, but the compilation model matches the repetition pattern.
• Compiler research. The project is small enough to inspect end to end. It exposes aliasing, state commits, cooperative occupancy, tensor-core tails, numerical gates, and the point where global fusion stops helping.

Three conditions recur: the model is small, the shapes are static, and the same compiled step runs enough times to repay setup. If any one of those is false, the case weakens quickly.

Where it is not useful#

TrainKernel is not a general training compiler. Version 0.2.0 does not implement attention, convolutions, dropout or random-number generation, dynamic shapes, distributed collectives, arbitrary input mutation, asynchronous task scheduling, or AdamW. Its optimizer story is explicit SGD inside a functional step. A different shape, stride, device, or dtype requires another compilation.

It is also not a credible path for conventional large-model training in its current form. The small and medium benchmark losses show why. Once matrix multiplication is large enough, specialized GEMM schedules and libraries have room to dominate. Holding the whole graph inside one cooperative grid becomes a constraint rather than a gift.

The measurements are from one A100 SXM4 80 GB and one H100 SXM 80 GB environment. They do not establish behavior across every SM80-or-newer GPU, long-run convergence, data-loader throughput, host-to-device transfers, multi-GPU scaling, or end-to-end application latency. GPU clocks could not be fixed in the recorded cloud environments. The reports retain this provenance rather than presenting the medians as hardware-independent constants.7TrainKernel, A100 optimization evidence, H100 optimization evidence, and GPU access record.

The benchmark's fastest-Inductor column includes every measured Inductor variant, even when a variant misses the declared numerical tolerance; each raw result separately records the fastest passing reference. That choice is conservative for the headline comparison. It prevents TrainKernel from winning because a faster baseline result was hidden, while keeping numerical failures visible for anyone deciding which configuration is usable.

What I learned#

The project began with a clean systems hypothesis: training contains more fusible structure than a sequence of forward operators, and very small training loops may benefit from collapsing that structure into one device program. The A100 and H100 results support that hypothesis in a real but narrow regime. TrainKernel wins the two smallest held-out shapes on both architectures, as well as the independent classification-head check on each, and it loses decisively as matrix work grows.

The next architectural step would not be “more fusion.” It would be a cost model. A practical compiler should recognize when a launch-bound graph belongs in a megakernel, when a GEMM should be delegated to Inductor or cuBLAS, and where to split a graph so both sides keep their advantage. The shape-aware dispatcher in the roadmap matters more than adding another GPU row.

I also came away more skeptical of instruction-level stories. The H100 WGMMA path is correct. Its isolated probe can be faster. Its training step is slower. The program around an instruction decides whether the instruction matters.

TrainKernel 0.2.0 is available under Apache 2.0 on GitHub and as a Python package. The repository includes generated CUDA, raw benchmark samples, correctness logs, rejected schedules, environment records, and reproduction commands. The most important output is not a universal claim that megakernels are faster. It is a measured boundary: for tiny repeated training, one kernel can be enough; beyond that boundary, the compiler has to know when to stop.

Source and documentation
A100 evidence · H100 evidence · Chart data
Original article text © 2026 Arya Somu, CC BY 4.0.