Graph Partitioning Explained

In Part 1 we opened an attention block and found 49 primitive operations: FULLY_CONNECTED, BATCH_MATMUL, SOFTMAX, and also SIN, COS, SELECT_V2, SLICE, TRANSPOSE.

Now the awkward question. You have an NPU. It has kernels for convolution and matrix multiply, because those are 90% of the FLOPs in any model and that’s what the silicon was designed around. Does it have a kernel for SIN?

Maybe. Maybe not. And if not, what happens to the model?

Partitioning

The graph gets cut into pieces, and each piece goes to whatever can run it:

[FC][FC][FC]  [SIN][COS]  [MATMUL][SOFTMAX][MATMUL][FC]
└─ subgraph A ─┘└─ B ──┘  └────── subgraph C ─────────┘
     NPU          CPU               NPU

This is graph partitioning, the step that turns “here is a model and here is some hardware” into an actual execution plan.

Two things here are worth internalizing before we go further, and neither one is obvious.

Every Cut Costs Something

The naive mental model is that partitioning is free bookkeeping: sort ops into buckets, run each bucket on the right device. It isn’t free bookkeeping, rather, the pieces don’t share memory.

Subgraph A’s output lives in NPU memory, in whatever layout the NPU wanted. The CPU can’t necessarily read that directly. So before SIN can run you may need a copy, a layout change, or both, and then the result has to go back.

Which means a graph cut into two pieces and the same graph cut into twenty pieces behave completely differently. In the pathological case, a heavily fragmented graph runs slower than pure CPU. You pay all the transfer overhead and never accumulate enough work on the accelerator to earn it back.

Anyone who has shipped an accelerated pipeline knows this failure mode: the NPU is busy, the profiler looks reasonable, and the end-to-end number is worse than the CPU baseline you were trying to beat.

It’s also why the buffer negotiation inside the LiteRT runtime is more than plumbing. The dispatch layer makes each accelerator declare which buffer types it accepts, and the docs are explicit that when two accelerators don’t overlap, the runtime inserts a conversion. From DISPATCH_API.md:

User -> GPU:  what buffers do you accept?
GPU  -> User: GlTexture, ClBuffer
User -> NPU:  what buffers do you accept?
NPU  -> User: AHWB
              Buffer conversion might happen (GlTexture -> AHWB)

“Zero-copy pipeline” is a property of a path, not of a component. It only holds when every stage’s accepted buffer types actually intersect. Chain a GPU preprocessing step into an NPU inference step whose buffer types don’t overlap, and you’ve quietly put a copy back into every frame.

Who Decides, and When

The second non-obvious thing: partitioning is not the vendor’s job.

JIT_COMPILATION.md describes the flow, and the wording is precise. The compiler plugin “compiles the partitioned subgraphs.” The framework does the cutting, and the vendor plugin receives subgraphs that are already carved out, then compiles them for its hardware.

That’s a good separation. A silicon vendor shouldn’t have to write a graph partitioner, they should only have to answer “do I support this op, and here’s the compiled code for the ones I do.”

When the cutting and compiling happen is a separate axis, and LiteRT supports both:

The official docs list five NPU vendors and not all support both. As of writing, Google Tensor is AOT-only in beta, while Qualcomm, MediaTek, Intel and Samsung support both paths.

The JIT Detail That’s Easy to Skim Past

It’s stated plainly in the LiteRT docs, and I still nearly went right over it:

Since JIT compilation relies on in-memory handles that exist only during the lifetime of the compiler plugin and the runtime process, JIT compilation cannot be cached.

If the runtime detects JIT handles it disables model caching for that run and logs JIT execution handles detected. Disabling JIT model caching., and this happens even if you configured a compilation cache directory. It gets bypassed.

The docs describe JIT’s downside as “a higher first-run cost.” Accurate, but easy to read too kindly. On a phone or a set-top box an app gets launched, backgrounded, reclaimed under memory pressure, then launched again. If nothing survives across process lifetimes, that “first run” cost is a cost you pay on every run.

So the AOT/JIT trade is not simply “AOT needs a toolchain, JIT doesn’t, therefore JIT is more convenient”:

On a memory-constrained device where processes get killed and restarted often, that last row can dominate the other two.

What This Changes in Practice

Three things I’ve started doing differently.

When inference is slower than expected, I now ask how many subgraphs there are before touching anything else. Fragmentation is a more common cause than a slow kernel, and it stays invisible unless you go looking for it. The runtime ships a profiler, and model-explorer will draw the graph for you.

The vendor’s op coverage list belongs in model selection, not in a footnote at the end of an evaluation. If you know which ops cause a break, you can pick or adjust a model to route around them, which beats tuning after the fact by a wide margin. Recall from Part 1 that RoPE was nine primitive ops, so how a vendor handles that pattern matters more than its headline TOPS number.

And when evaluating a chip, ask which buffer types its NPU accepts, not just whether it has an NPU. If those types don’t intersect with your camera or GPU preprocessing path, you’ll be paying a conversion on every frame and no amount of TOPS fixes that.

Part 3 goes one layer down, into the Dispatch API: the actual C interface a silicon vendor implements to plug an NPU into LiteRT, why it replaced the TFLite Delegate, and what that interface tells you to ask a vendor before you commit to their silicon.

Sources are the design docs in google-ai-edge/LiteRT at commit dc32e93, JIT_COMPILATION.md, DISPATCH_API.md and COMPILER_PLUGIN.md, plus the NPU acceleration page on the Google AI Edge developer site.