Kernel fusion

Code written with a library like PyTorch reads one line at a time: add these two arrays, multiply by that, take the exponential. Run as it is written, each of those lines becomes its own GPU kernel. Each kernel reads its inputs from the GPU’s memory, does a tiny amount of arithmetic and writes its result back, and every launch has a fixed start-up cost of a few microseconds.

When the arrays are small, that start-up cost is most of the time. A line like

$r = e^{-40\,\overline{(q - q_{\text{ref}})^2}}$

is four kernels (subtract, square, average, exponential), each finishing in a blink and then waiting for the next one to start.

Kernel fusion merges a chain like that into one kernel that reads $q$ and $q_{\text{ref}}$ once, does all four operations in the GPU’s fast registers, and writes $r$ once. Fewer launches, fewer trips to memory.

torch.compile does it automatically. It traces the Python code into a graph of operations, works out which ones can be merged, and writes new fused kernels (in a language called Triton) for the GPU. That takes minutes the first time, so it pays off for code that runs millions of times, like a training loop. Code it can’t trace, such as calls into another GPU library, is left as it is, and the compiled pieces run around it.

Fusion and a CUDA graph attack the same waste from two sides: the graph removes the CPU’s cost of sending each kernel, fusion removes most of the kernels.

In this project

In Two fighters and vision, one step of the opponent environment ran about 13,700 kernels for 2,048 copies of the world: about 1,800 for the MuJoCo Warp physics and 11,700 small ones for the rewards, observations and resets written in PyTorch. Compiling that PyTorch code (the physics calls left out) cut a step on an RTX 5090 from 79.5 to 60.2 ms, and from 58.7 to 41.7 ms with 1,024 copies. It costs about four minutes of compiling at the start of each run.