Kernwerk

Compilation: one engine instead of a thousand launches

A trained model is usually slower than it needs to be, not because of the maths, but because of how that maths gets run.

Take PyTorch, the framework almost everyone trains in. It runs a model one layer at a time. For each layer a separate instruction sent to the GPU from Python, one after another. That flexibility is ideal while you experiment, but at inference it means thousands of tiny hand-offs, and the GPU spends much of its time waiting for the next instruction instead of computing.

A compiler removes that waiting. It takes the finished model and, ahead of time, turns it into a single optimised program built for the exact chip it will run on. The weights are untouched and the accuracy is unchanged: the model is simply run properly.

What the compiler does

We used TensorRT, NVIDIA’s inference compiler. “Compiler” is shorthand for four jobs:

  • It fuses layers. Steps that were separate (a convolution, a normalisation, an activation) collapse into one. Fewer instructions, fewer trips to memory.
  • It tunes for the chip. Each layer has several possible implementations. At build time TensorRT runs them on the target GPU and keeps the fastest. It doesn’t guess from a rulebook. It measures. That is why the build takes minutes, and why the result only fits one kind of chip.
  • It plans memory. Layers reuse the same working memory instead of each reserving its own, so the model needs far less of it.
  • It runs the model itself. The output is a single compiled engine that executes end to end on the GPU, with no return to Python between layers.

The result

To make it concrete: our example is a chest-X-ray model, a DenseNet-121 that reads an X-ray image and scores it for 14 possible diseases, running on the same $249 Jetson Orin Nano as in Scan in, report out. We run it two ways, both in 16-bit precision, both taking 8 images at a time, so the only thing that differs is how it is executed:

  • PyTorch: run layer by layer, the way you train it.
  • Compiled (TensorRT): the same model turned into a single engine.
Execution overview: PyTorch dispatches and waits for each layer through Python; the compiled engine runs the whole model in a single call.

Compiled with TensorRT, the model runs close to three times faster than it does in PyTorch, and the accuracy does not move: a macro-AUROC of 0.7405 either way, measured on 2000 held-out chest X-rays (roughly, how reliably it ranks a diseased scan above a healthy one). Nothing about the model itself changed, we just stopped running it wastefully.

Throughput Speedup macro-AUROC
PyTorch, 16-bit 169.5 img/s 0.7405
Compiled (TensorRT), 16-bit 499.1 img/s 2.9× 0.7405

Both rows are the same weights at the same precision on the same images, so that 2.9× is compilation and nothing else. Each throughput is the mean of three runs of 400 images, with under 1 img/s of spread between runs, and the AUROC comes from the same 2000 scans in both cases.

The boundary

Compilation is not free. An engine is tied to one chip and one version of the tool: build it on our Jetson and it won’t run anywhere else. You rebuild it for every model, every precision, every chip.

And the tool doesn’t travel. TensorRT is NVIDIA-only. Everywhere else you get the vendor’s own stack (NXP eIQ, TI TIDL, Rockchip RKNN, Espressif ESP-DL), and those are thinner. One operator the accelerator doesn’t support sends the tensor back to the CPU and back again, on every frame.

The upside is that the headroom is much larger than on a Jetson. On a $5 ESP32 we took a model from 9.3 seconds to 92 milliseconds per inference.

That’s where we come in: the right toolchain for every chip you ship on, rewritten where it falls short, with the accuracy measured before and after.