ML compiler stacks for edge inference: graph lowering, operator fusion and kernel selection across TensorRT, TVM, ONNX Runtime and TFLite

On a datacenter GPU, a model is mostly its weights. On an edge device, a model is its weights plus whatever the compiler did to them. The same ONNX file can run several times faster, or silently lose accuracy, depending on how it was lowered, which operators were fused, which kernels were picked and which parts of the graph quietly fell back to the CPU. If you ship models to Jetson modules, mobile SoCs, x86 gateways or NPU-equipped cameras, the compiler stack is part of your model, and it needs the same engineering discipline as training.

This post walks through what actually happens between a trained network and an executable engine on edge silicon, how the four stacks we see most often (TensorRT, Apache TVM, ONNX Runtime and TFLite/LiteRT) differ, where things break, and how to benchmark the result honestly.

The lowering pipeline, stage by stage

Every ML compiler does some version of the same four stages. Knowing which stage owns a problem is most of the debugging work.

1. Capture and export. The framework graph is traced or exported (for example via torch.export or the ONNX exporter) into a static IR. Decisions made here are sticky: opset version, which dimensions are dynamic, whether control flow is unrolled, and how custom ops are represented. A model exported with an unnecessary dynamic batch axis can lose most of the downstream optimisation opportunities.
2. Graph-level passes. Constant folding, shape inference, dead-code elimination, common-subexpression elimination, BatchNorm folding into preceding convolutions, and layout transformation (NCHW to NHWC or a vendor-specific blocked layout). Layout is one of the biggest silent wins and losses: a single op that only supports the "wrong" layout can force transpose pairs around it.
3. Operator-level transformation. Fusion, tiling, vectorisation and precision assignment. This is where the graph stops looking like your PyTorch module.
4. Kernel selection and code generation. The compiler either picks from a library of hand-tuned kernels ("tactics"), generates code from a schedule, or offloads subgraphs to a vendor delegate. The result is specific to the target architecture and, often, to the exact compiler version.

Operator fusion: why it matters and when it hurts

Most edge workloads are memory-bandwidth bound, not compute bound. A useful mental model is the roofline: if an operator's arithmetic intensity (FLOPs per byte moved) sits below the device's ridge point, it will be limited by DRAM bandwidth no matter how fast the tensor cores are. Elementwise ops such as bias-add, activation, residual add and scaling have very low intensity. Running them as separate kernels means writing an intermediate tensor to memory and reading it back each time, plus a kernel launch.

Fusion removes those round trips. The common patterns:

• Epilogue (vertical) fusion: Conv/GEMM followed by bias, BatchNorm (already folded), activation and residual add executed in one kernel, with the epilogue applied while the output tile is still in registers or shared memory.
• Horizontal fusion: several independent ops on the same input (for example the Q, K and V projections) merged into one wider GEMM.
• Pattern fusion of whole blocks: attention (QKT, scale, mask, softmax, V-multiply) computed tile-by-tile without materialising the full attention matrix, in the style popularised by FlashAttention; LayerNorm and GELU fused into neighbouring GEMMs.

Fusion is not free. Large fused kernels increase register pressure and can reduce occupancy. Fusing across a quantise/dequantise boundary can change where rounding happens, which changes numerics. Some fusions recompute values instead of storing them, which is a win on bandwidth-starved devices and a loss on compute-starved ones. And a fused kernel is one opaque node, so per-layer debugging becomes harder. A compiler that fuses aggressively still needs to tell you what it fused.

How the four stacks differ

TensorRT

TensorRT parses ONNX (or a network built through its API), applies graph optimisations and fusions, then times candidate kernels on the actual GPU during the build to pick tactics. The output engine is tied to the GPU architecture and the TensorRT version, so it is a build artifact per target, not a portable model. Practical levers: precision flags (FP16, INT8, FP8 on newer architectures), optimisation profiles for dynamic shapes (min/opt/max), workspace limits, timing caches to make builds faster and more repeatable, and plugins for unsupported ops. Because tactic selection is empirical, two builds on a noisy or thermally throttled device can choose different kernels.

Apache TVM

TVM lowers through its own graph IR (Relay, and its successor Relax) into TensorIR, then searches for good schedules with an auto-tuner (MetaSchedule). Its strength is reach: it can generate code for CPUs, GPUs and accelerators that have no mature vendor runtime, and its "bring your own codegen" path offloads supported subgraphs to vendor libraries while compiling the rest itself. The cost is engineering time: tuning should run on the real target (typically via an RPC tracker to devices on a bench), tuning logs must be versioned, and the search budget is a real line item.

ONNX Runtime

ONNX Runtime is a runtime first and a compiler second. It applies graph optimisations at configurable levels, then partitions the graph across execution providers: CUDA, TensorRT, OpenVINO, QNN, XNNPACK, NNAPI, CoreML and others, with CPU as the universal fallback. That makes it very good for portability across a heterogeneous fleet. The trap is partitioning: every boundary between providers is a potential device-to-host copy, and a graph split into many small islands can be slower than running entirely on CPU.

TFLite / LiteRT

The TFLite converter (now under the LiteRT name) produces a flatbuffer that runs on the built-in CPU kernels, with delegates for the GPU, NNAPI, Hexagon/QNN, XNNPACK and others. NPUs on mobile and embedded SoCs usually require full-integer quantisation and support a limited op set; anything unsupported is executed on the CPU, again creating partition boundaries.

Failure modes we design against

• Silent fallback. One unsupported op (a custom NMS, an unusual resize mode, a dynamic slice) drops a subgraph to CPU. The model still runs and still returns correct results, just much slower, and often only on one device SKU.
• Numerical divergence. FP16 accumulation, fused softmax and reordered reductions change outputs. Most layers tolerate it; LayerNorm variance, large-logit softmax and box-regression heads often do not. Overflow to inf in FP16 is a classic, and it may only appear on rare inputs.
• Dynamic shape penalties. Kernels tuned for one shape run poorly at others, or the runtime rebuilds/re-plans at first sight of a new shape and causes a latency spike in production.
• Non-reproducible builds. Empirical tactic selection plus thermal noise means two "identical" builds can differ in latency and, occasionally, in numerics.
• Version lock. An engine built with one JetPack/TensorRT version will not load on another. Fleet upgrades become model rebuilds.
• Pre/post-processing drift. Resize interpolation, colour order and normalisation done differently on device than in training. This is not a compiler bug, but it shows up at the same time and gets blamed on the compiler.

Benchmarking methodology that survives scrutiny

Edge benchmark numbers are easy to inflate by accident. We hold ourselves to a few rules:

• Pin the device state. Record power mode and clock settings (for example the nvpmodel mode on Jetson), ambient conditions and firmware versions. Compare like with like.
• Separate warm-up from steady state, and run long enough to hit thermal equilibrium. A one-minute benchmark on a fanless enclosure measures the heatsink's thermal mass, not sustained performance.
• Report distributions, p50/p95/p99 latency and sustained throughput, at the batch size and concurrency the product will actually use.
• Measure end to end (decode, preprocess, inference, postprocess, output) as well as engine-only time. Fusion that saves 2 ms in the engine is irrelevant if preprocessing costs 15 ms on the CPU.
• Measure power and memory, not just speed. On shared-memory SoCs, engine workspace competes with video buffers.
• Gate on correctness twice: output-level comparison against a reference runtime with explicit tolerances, and task-level metrics on a held-out set that includes the hard tail.

For output-level comparison, Polygraphy and trtexec make the loop cheap:

# compare TensorRT FP16 against ONNX Runtime on the same inputs
polygraphy run model.onnx --trt --fp16 --onnxrt \
  --atol 1e-3 --rtol 1e-3

# build and time an engine reproducibly with a shared timing cache
trtexec --onnx=model.onnx --fp16 \
  --timingCacheFile=cache/orin.tcache \
  --saveEngine=build/model_orin_fp16.plan

When outputs diverge, bisect: mark intermediate tensors as outputs, find the first layer that exceeds tolerance, then force that layer back to higher precision or exclude it from fusion.

Deployment patterns

• Engines are build artifacts, not models. Ship a manifest with each one: source model hash, compiler and runtime versions, target architecture, precision per layer, input shapes and profiles, calibration-set hash and the correctness report.
• Build per SKU in CI, on real hardware or an exactly matching module, with timing caches checked in so builds are faster and more repeatable.
• Keep a portable fallback (for example an ONNX Runtime CPU path) so a device that cannot load its optimised engine degrades rather than going dark.
• Treat runtime upgrades as model releases, with the same evaluation and staged rollout.

How HOLORAI approaches this

We start from the device and the budget, not the model. Before choosing a stack we profile the target: memory bandwidth, supported precisions, accelerator op coverage and the thermal envelope of the actual enclosure. That usually decides the compiler more than any benchmark chart does: TensorRT where the fleet is NVIDIA and homogeneous, ONNX Runtime with selected execution providers where hardware is mixed, TFLite/LiteRT delegates on mobile and NPU SoCs, and TVM where a target has no mature vendor runtime.

We then make the model compiler-friendly at export time: static shapes where the product allows, supported resize and NMS variants, and custom ops replaced or implemented as plugins with their own tests. Every build produces an engine, a manifest and a correctness report, and the pipeline fails if any subgraph falls back to an unexpected device or if any output exceeds its tolerance. Latency, power and memory are measured end to end at thermal steady state, on the same enclosure the customer will deploy.

The goal is not the fastest number on a slide. It is an engine whose speed, accuracy and failure behaviour we can explain, reproduce and rebuild when the platform changes underneath it.