LiteRT Exploration - Analyzing attention.tflite

I was reading the LiteRT runtime source to understand how it decides which parts of a model go to the NPU and which fall back to CPU. Somewhere in the middle I realized I had skipped a step. Before we can talk about splitting a model, we need a clear answer to something more basic: what is a model, physically, as a file?

Short version: a model is a pile of numbers plus a set of instructions for what to do with them. The numbers are the weights. The instructions are a graph. Almost all the bytes are weights, and almost all the interesting engineering lives in the graph.

Let’s crack open a real .tflite and look at both.

Why a Graph Falls Out Naturally

A neural network is function composition. That’s the whole thing:

y = f₄(f₃(f₂(f₁(x))))

Each f is a layer, a convolution, a matrix multiply, a normalization, an activation. Input flows through, gets reshaped and transformed, output comes out.

That “one after another” structure is a chain. Add skip connections (ResNet) or multi-input operations (attention needs Q, K and V arriving together) and the chain becomes a directed acyclic graph. Nodes are operations, edges are tensors flowing between them.

So the graph isn’t an abstraction someone invented to look clever. It’s a direct transcription of data dependencies that were always there.

Why Not Just Ship Code?

You wrote the model in PyTorch. That’s Python. Why not ship the Python?

Two reasons, and the second one is the real answer.

First, there’s no Python interpreter on the device. Fine, you could compile it.

Second: an NPU does not execute programs. It’s a block of silicon that does matrix multiplication very fast. You cannot hand it a for loop. You can hand it “here is a convolution, here are its weights, here is where the input lives.” The unit it understands is an operation, not an instruction stream.

So we need something in the middle that is portable, inspectable, and transformable. This is exactly what a compiler does:

C source  ->  IR  ->  x86 / ARM / RISC-V machine code

And the on-device ML version of the same idea:

PyTorch (Python)  ->  torch.export -> FX Graph  ->  LiteRT FlatBuffer  ->  CPU / GPU / NPU
    authoring              intermediate repr          the portable graph      execution

Google’s litert-torch is the middle arrow. It used to be called ai-edge-torch, and the rename is recent enough that most search results still land on the old name.

Ops Are an ABI

Here’s the framing that made everything else click for me.

An operation is a contract between two groups of people who never talk to each other.

The op set is effectively an ABI. Model authors compose ops, vendors implement ops, and neither side needs to know anything else about the other.

Which is why “which ops does this NPU support?” is the question that decides everything downstream. Hold that thought, it’s the whole subject of Part 2.

What’s Actually in the File

A .tflite is a FlatBuffer. Strip the serialization details and it holds four things:

Operators reference tensors by index, and those indices are the edges. The file is a serialized graph, quite literally.

Opening a Real One

I wrote a small dumper against the tflite Python bindings and pointed it at attention.tflite, a test model from the LiteRT-LM repo. It’s a single attention block exported from JAX.

file            : attention.tflite  (10,399,276 bytes on disk)
schema version  : 3
subgraphs       : 1
operator codes  : RESHAPE, FULLY_CONNECTED, CAST, DIV, SIN, COS, SLICE, MUL,
                  SUB, ADD, CONCATENATION, TRANSPOSE, BATCH_MATMUL,
                  SELECT_V2, SOFTMAX
weight bytes    : 10,387,852  (99.9% of the file)

99.9% of the file is weights. 78 tensors, 49 operators, and the entire description of what to compute fits in roughly 11 KB. By size the graph is rounding error. It is also the only part that determines whether the model runs on an accelerator at all.

The three inputs tell you what an attention block needs:

tensor 0   BOOL     8x1x100      serving_default_args_3:0    <- attention mask
tensor 1   FLOAT32  8x100x128    serving_default_args_0:0    <- hidden states
tensor 2   INT32    8x100        serving_default_args_1:0    <- token positions

Batch 8, sequence length 100, model dim 128.

You Can Read the Algorithm Off the Ops

This is the part I found genuinely fun. Here is the operator list with light annotation. None of it is my interpretation of source code, it’s just what sits in the file.

op[1]  FULLY_CONNECTED   in=[29, 7]   -> 30    ┐
op[3]  FULLY_CONNECTED   in=[29, 6]   -> 32    ├─ Q, K, V projections
op[5]  FULLY_CONNECTED   in=[29, 5]   -> 34    ┘  (same input 29, three weights)

op[8]  CAST              in=[36]      -> 37    ┐
op[9]  DIV               in=[37, 22]  -> 38    ├─ RoPE: build the angle table
op[11] SIN               in=[39]      -> 40    │   from token positions
op[12] COS               in=[39]      -> 41    ┘

op[13] SLICE   ┐
op[14] SLICE   │
op[15] MUL     │
op[16] MUL     ├─ rotate-half applied to Q
op[17] SUB     │   (x₁cos - x₂sin, x₂cos + x₁sin)
op[18] MUL     │
op[19] MUL     │
op[20] ADD     │
op[21] CONCAT  ┘
op[22]..op[30]  ── the identical nine-op pattern, applied to K

op[31] MUL               in=[50, 28]  -> 60    <- scale by 1/√d
op[38] BATCH_MATMUL      in=[66, 63]  -> 67    <- Q·Kᵀ, the attention scores
op[41] SELECT_V2         in=[69,68,27]-> 70    <- causal mask (that BOOL input)
op[42] SOFTMAX           in=[70]      -> 71
op[44] BATCH_MATMUL      in=[72, 65]  -> 73    <- attn · V
op[48] FULLY_CONNECTED   in=[76, 4]   -> 77    <- output projection

Scaled dot-product attention with rotary position embeddings, written out in 49 primitive operations. There is no attention op and no RoPE op here, just slices, multiplies, sines and cosines.

The Weight Shapes Leak the Architecture

Look at the three projection weights:

tensor 7   FLOAT32  128x128   65,536 bytes   <- Q projection
tensor 6   FLOAT32   16x128    8,192 bytes   <- K projection
tensor 5   FLOAT32   16x128    8,192 bytes   <- V projection

Q projects 128 -> 128. K and V project 128 -> 16. That asymmetry is not a typo.

Further down the graph, tensor 74 has shape 8x32x100x4: batch 8, 32 heads, sequence 100, head dim 4. So the query side is 32 × 4 = 128. On the key/value side, 16 ÷ 4 = 4 heads.

32 query heads sharing 4 key/value heads. That’s grouped-query attention with a group size of 8, and you can read it straight off the tensor shapes without opening any documentation. Fewer KV heads means a smaller KV cache, which is the trade GQA exists to make.

Why This Matters More Than It Looks

Three practical consequences, and they each become the subject of a later post.

Weights alone are not a model. The same 10 MB of floats could be convolution kernels or a plain matrix. The graph is the instruction manual and the weights are the parts, which is why you can’t just “load the weights” onto an NPU.

The op list is the compatibility surface. Look at it again: SIN, COS, SELECT_V2, BATCH_MATMUL. Every one of those is a question mark for an NPU. A chip that does convolutions beautifully might have no SIN kernel at all, and when that happens the graph gets cut into pieces, some on NPU and some on CPU, with a cost at every cut.

Fusion happens at this level. Because RoPE is spelled out as nine primitive ops rather than one, a compiler can pattern-match those nine and swap in a single fused hardware primitive, assuming the vendor has one. That’s why the LiteRT repo ships a PATTERN_MATCHING.md.

Part 2 looks at what happens when a chip only runs part of that op list: how the graph gets partitioned, who decides the cut, when the decision happens, and why cutting badly ends up slower than not using the accelerator at all.

The dump tool is about 60 lines of Python against the tflite package. The model is schema/testdata/attention.tflite in google-ai-edge/LiteRT-LM, read at commit effe245.