What I Learned Building a Mini TensorRT

Anushka S Raghunandan — Mon, 07 Sep 2026 19:15:23 +0000

I wanted to understand how inference compilers like NVIDIA's TensorRT actually make models faster — not the marketing explanation, the mechanical one. So instead of reading the docs and calling it a day, I built a small version myself: OptiFlow, a graph compiler that takes an ONNX model and optimizes it the way a production inference engine would.

Here's what it does, why it works, and the bug that taught me the most.

The idea: models are graphs, and graphs can be rewritten

An ONNX model is just a computation graph — nodes are operators (Conv, Add, Relu, etc.), edges are tensors flowing between them. Most people treat that graph as fixed once training is done. It isn't. OptiFlow lifts the ONNX graph into a custom intermediate representation (IR) and runs a series of optimization passes over it:

Constant folding — if a subgraph's inputs are all constants, compute it once at compile time instead of on every inference call.
Dead-node elimination — remove nodes whose outputs are never consumed.
Operator fusion — merge adjacent operators (Conv → Relu, Add → Relu) into a single fused op. The first two are intuitive. Fusion is where the interesting engineering is.

Why fusion actually matters

Naively, Conv → Relu looks like it should cost roughly the same whether you run it as two operators or one. It doesn't, and the reason is memory, not compute.

Run them separately, and the hardware has to: compute the Conv output, write the full tensor back to memory, then read that same tensor back in for Relu, before writing the final result out again. That's two extra memory round-trips for a value that's only ever used once. On modern accelerators, moving data is frequently more expensive than the arithmetic itself — kernels are often memory-bandwidth-bound, not compute-bound. Fusing the two operators into one means the intermediate Conv output never leaves fast on-chip memory before Relu consumes it — it disappears entirely.

This is also why fusion passes are a core part of what real inference compilers do: on GPUs, avoiding redundant memory traffic tends to matter more than shaving arithmetic.

On a ResNet-18 graph, OptiFlow's fusion pass collapsed 51 nodes down to 34 (17 fused operator pairs), which brought estimated memory traffic down about 19% and modeled p95 tail latency down about 27%.

The part that scared me: how do I know it's still correct?

This is the part nobody warns you about when you start writing a compiler pass: it's trivially easy to write a fusion rule that's structurally elegant and semantically wrong. A graph rewrite that looks correct on paper can silently change the model's output, and you won't know until inference numbers quietly drift.

So before trusting any transformation, OptiFlow runs it through a round-trip harness: after each optimization pass, the modified IR is serialized back into a valid ONNX graph, and both the original and the optimized graphs are run through onnxruntime on identical inputs. The outputs are diffed within floating-point tolerance. If they don't match, the pass is rejected — full stop, no partial credit.

This caught more bugs than I expected. The most instructive one: a fusion rule that merged an Add → Relu pair without checking whether the intermediate Add output had more than one consumer downstream. Most of the time it only fed into Relu, so fusing was safe. But in a branch of the graph where the Add output also fed a skip connection, fusing silently dropped that second consumer — the fused node produced the right final tensor for the Relu path, but the skip connection lost its input entirely. Structurally the graph still "looked" valid; numerically it was wrong. The round-trip harness caught the mismatch immediately, which is exactly the point of building it before optimizing anything else.

What I'd take into a real inference-compiler team

The lesson wasn't really about ONNX or fusion rules specifically — it was that optimization and correctness have to be verified together, not sequentially. It's tempting to write a rewrite pass, eyeball a few examples, and assume it generalizes. It's much harder — and much more valuable — to build the harness that proves it, before you trust the speedup.

If you're curious, the code is on GitHub — feedback and PRs welcome.

DEV Community: Anushka S Raghunandan

What I Learned Building a Mini TensorRT

The idea: models are graphs, and graphs can be rewritten

Why fusion actually matters

The part that scared me: how do I know it's still correct?

What I'd take into a real inference-compiler team