# Native Model Authoring for Edge with LiteRT Tensor API

> Source: <https://developers.googleblog.com/native-model-authoring-for-edge-with-litert-tensor-api/>
> Published: 2026-10-19 17:18:05.482423+00:00

Shipping on-device AI means meeting tight latency and memory budgets on every device you target. The most common path for deploying models in production is converting a trained model, for example converting a PyTorch model with [LiteRT Torch](https://github.com/google-ai-edge/litert-torch), and running it with the [LiteRT runtime](https://developers.google.com/edge/litert/inference). This offers the quickest and most reliable route to production, and it's the path we recommend for most projects. Sometimes, though, developers want full control over how a network is represented and optimized to support new architectures or to try a new lowering logic without waiting for upstream framework updates.

For ultimate control, developers write computations manually in native C++ or shader languages. [Apple's MLX](https://github.com/ml-explore/mlx) made this code-first approach popular on Apple silicon, and [llama.cpp](https://github.com/ggml-org/llama.cpp) showed how far tailored C++ implementations can take LLM inference. However, the level of control comes at a significant cost: implementations are tightly coupled to specific backends (such as Metal, OpenCL, or Vulkan) and must be rewritten for each new hardware target. Furthermore, managing zero-copy buffer interoperability and synchronization across CPU, GPU, and NPU, while gluing together pre-processing, inference, and post-processing, demands substantial manual engineering effort.

To solve the above challenges, we are excited to announce the release of the novel [LiteRT Tensor API](https://github.com/google-ai-edge/LiteRT/blob/main/tensor/README.md), a lightweight authoring layer in the LiteRT stack, designed to improve handcrafting models. It combines the flexibility of native authoring with cross-platform portability. You build the computation graph op by op in C++, choosing exactly which ops run and how they connect. The same graph can run through LiteRT on CPU, GPU (powered by open-source [ML Drift](http://github.com/google-ai-edge/ml-drift)) and NPU, or without model file or runtime directly on CPU through XNNPACK, so you don't rewrite it for each new hardware target. Model inference and pre/post-processing can be assembled with the same API and chained through shared accelerator buffers, so that data never round-trips through the CPU. Start with the conversion path, which remains the best default for most models. Reach for the Tensor API when you need low-level control over the graph or want the whole pipeline to run end to end on the accelerator.

The LiteRT Tensor API also serves as an **ideal programmatic foundation for AI coding agents,** providing structured primitives to generate, rewire, and tune model graphs automatically. In the rest of this post, we’ll dive into the design and walk through practical code examples showing how to author models from scratch.

To show the Tensor API in practice, we built an interactive [Segment Anything App](https://storage.googleapis.com/litertjs_demos/sam2_web/index.html), running entirely on-device in the browser on WebGPU. The app enables real-time multi-object segmentation and tracking via Meta's [Segment Anything Model 2 (SAM 2)](https://github.com/facebookresearch/sam2), triggered either by simple clicks or natural language prompts powered by [Gemma 4 E2B](https://huggingface.co/litert-community/gemma-4-E2B-it-litert-lm).

The SAM 2 model is fully authored in C++ via the LiteRT Tensor API. A coding agent wrote and evolved it from under 10 to over 30 frames per second at 512x512 on WebGPU. See that the complete implementation and the web app are open source in the [LiteRT Samples](https://github.com/google-ai-edge/litert-samples/tree/main/models/sam2/sam2_hiera_tiny_video) GitHub repo.

The LiteRT Tensor API provides four key features for authoring and optimizing on-device models:

We ran the same C++ code for the SAM 2 model on the GPUs of three devices: a macOS laptop (Metal), an iOS phone (Metal), and an Android phone (OpenCL). In our tests on the MacBook Pro, it ran about 1.5x faster than Apple MLX (fp16, image encoder plus mask decoder). The SAM 2.1 [sample](https://github.com/google-ai-edge/litert-samples/blob/main/models/sam2/sam2_hiera_tiny_video/tensor_api/sam2_image/README.md) includes the benchmark commands so you can reproduce the results.

With the Tensor API, each layer is a C++ function over [Tensor](https://github.com/google-ai-edge/LiteRT/blob/6f622b5ddd51429c564f5ac2b22a7861bedbd147/tensor/tensor.h#L159) values, built from high-level operators and weights loaded directly from the checkpoint. For example, here is a multi-head self-attention layer:

```
// Multi-head self-attention over x: [batch, tokens, dim].
Tensor Attention(const Tensor& x, const WeightMap& w, int heads) {
  std::vector<int> s = x.GetShape();
  int b = s[0], n = s[1], dim = s[2], d = dim / heads;
  Tensor q = FullyConnected(x, w.at("q_proj.weight"), w.at("q_proj.bias"));
  Tensor k = FullyConnected(x, w.at("k_proj.weight"), w.at("k_proj.bias"));
  Tensor v = FullyConnected(x, w.at("v_proj.weight"), w.at("v_proj.bias"));
  // Split heads: [b, n, dim] -> [b, heads, n, d].
  q = Transpose(Reshape(q, {b, n, heads, d}), {0, 2, 1, 3});
  k = Transpose(Reshape(k, {b, n, heads, d}), {0, 2, 1, 3});
  v = Transpose(Reshape(v, {b, n, heads, d}), {0, 2, 1, 3});
  // softmax(q * k^T / sqrt(d)) * v, for all heads at once.
  Tensor scores = BatchMatMul(q, k, /*adj_x=*/false, /*adj_y=*/true);
  Tensor attn = Softmax(Mul(scores, 1.0f / std::sqrt(static_cast<float>(d))));
  Tensor out = BatchMatMul(attn, v);
  // Merge heads: [b, heads, n, d] -> [b, n, dim], then project.
  out = Reshape(Transpose(out, {0, 2, 1, 3}), {b, n, dim});
  return FullyConnected(out, w.at("out_proj.weight"), w.at("out_proj.bias"));
}
```

In the SAM 2 sample, [BuildEncoder](https://github.com/google-ai-edge/litert-samples/blob/fe65fa5512b273bcd62936afd18b924af7d5ae2c/models/sam2/sam2_hiera_tiny_video/tensor_api/sam2_image/sam2_graph.cc#L765-L833) and [BuildDecoder](https://github.com/google-ai-edge/litert-samples/blob/fe65fa5512b273bcd62936afd18b924af7d5ae2c/models/sam2/sam2_hiera_tiny_video/tensor_api/sam2_image/sam2_graph.cc#L835-L1006) compose layers like this into the encoder and decoder graphs, and [ModelFactory](https://github.com/google-ai-edge/LiteRT/blob/6f622b5ddd51429c564f5ac2b22a7861bedbd147/tensor/backends/tflite/tflite_flatbuffer_conversion.h#L47) exports both as two signatures of one `sam2_image.tflite` file.

A `.tflite` file is one of two ways to run your graph. The backend is a template argument on `Tensor: TfLiteMixinTag` exports a `.tflite` that LiteRT runs on CPU, [ML Drift](https://developers.googleblog.com/ml-drift-next-gen-gpu-aiml-inference-at-the-edge/) (GPU), or NPU, and `XnnpackMixinTag` runs the same code directly on CPU through XNNPACK, with no model file. Tag tensors with both to keep both options:

```
// Tag tensors for both backends; ops record one graph either can run.
using T = Tensor<TfLiteMixinTag, XnnpackMixinTag>;
T y = Attention(x, w, /*heads=*/8);  // The layer above, over T values.
// Export a .tflite and run it with LiteRT on CPU, GPU, or NPU.
Save({y}, "attention.tflite");
// Or run the same graph directly on the CPU with XNNPACK.
auto runner = XnnpackRunner::Create({y});
runner->Run();
```

SAM 2 takes the `.tflite` path to reach WebGPU in the browser, and the rest of this post follows it. The [Gemma 4 example](https://github.com/google-ai-edge/LiteRT/tree/main/tensor/examples/gemma4) takes the direct path for prefill and decode on CPU, the lightest way to run Tensor API code. The [Tensor API README](https://github.com/google-ai-edge/LiteRT/blob/main/tensor/README.md) walks through both.

A natural workflow uses both: iterate on the direct XNNPACK path, loading weights straight from the `.safetensors` checkpoint with no model file to rebuild, then export a `.tflite` when you're ready to ship with LiteRT on mobile or on the web through LiteRT.js.

A real app needs more than the neural network itself. In the demo, turning a camera frame and a click into a mask takes four steps: pre-processing converts the frame into the 512x512 image the model expects; the image encoder turns that image into features; the mask decoder combines those features with the click to predict a mask; and post-processing scales the mask back to the frame's size. Running this whole pipeline fast takes two things: keeping data on the accelerator (such as the GPU) between steps, and controlling how each step runs on each accelerator.

In a typical app, only the model runs on the GPU; the CPU resizes each camera frame, passes the encoder's output to the decoder, and post-processes the mask. Each trip copies data and stalls the GPU.

In [litert-samples/sam2](https://github.com/google-ai-edge/litert-samples/tree/main/models/sam2/sam2_hiera_tiny_video/tensor_api), all four steps are Tensor API graphs on the GPU. Pre-processing, for example, is a small graph built with [CreateLambdaRunner](https://github.com/google-ai-edge/LiteRT/blob/6f622b5ddd51429c564f5ac2b22a7861bedbd147/tensor/runners/litert/lambda_model_runner.h#L143). The graphs are connected once at startup to share buffers, then run for each frame and click:

```
// Pre-processing is authored with the same Tensor API ops as the model:
// RGBA camera frame [1, H, W, 4] -> normalized pixels [1, 512, 512, 3].
auto pre_runner = CreateLambdaRunner(
    env, options,
    {{"frame", Tensor({.type = Type::kU8, .shape = {1, H, W, 4}})}},
    [=](const TensorsMap& inputs) {
      Tensor rgb = Slice(inputs.at("frame"), {0, 0, 0, 0}, {1, H, W, 3});
      Tensor resized = ResizeBilinear(Cast(rgb, Type::kFP32), {512, 512});
      return TensorsMap{{"pixels", Add(Mul(resized, scale), bias)}};
    });
auto model = LitertDynamicRunner::Create(env, "sam2_image.tflite", options);
// Connect once: neighboring steps share buffers, so nothing is copied.
pre_runner.SetOutput("pixels", model.GetInput("encode_image", "pixels"));
for (auto name : {"image_embeddings", "feat_s1", "feat_s0"}) {
  model.SetInput("decode_mask", name, model.GetOutput("encode_image", name));
}
post_runner.SetInput("masks", model.GetOutput("decode_mask", "masks"));
```

No data returns to the CPU between steps: all 522 encoder and 269 decoder operations run on the GPU, and pre- and post-processing add only about 0.34 ms per frame.

Because the graph is plain C++, you or your coding agent can rewrite it for each backend. Our agent, [Antigravity](https://antigravity.google/), moved SAM 2's per-frame decisions, such as whether the object is still visible, from CPU code into the graph. Comparisons like score > 0 would still fall back to the CPU on WebGPU, because they output true/false values, so the agent rewrote them as arithmetic in [StepPos](https://github.com/google-ai-edge/litert-samples/blob/fe65fa5512b273bcd62936afd18b924af7d5ae2c/models/sam2/sam2_hiera_tiny_video/tensor_api/sam2_video/sam2v_graph.cc#L191-L198):

```
// Simplified from StepPos and BuildStep in sam2v_graph.cc.
// Before: Greater outputs true/false values, which fall back to the CPU.
// Tensor visible = Cast(Greater(score, zero), Type::kFP32);
// After: GPU-friendly ops; 0 for score <= 0, 1 for score >= 1e-4.
Tensor visible = Sub(1.0f, Relu(Sub(1.0f, Mul(Relu(score), 1e4f))));
// Keep the mask if the object is visible; otherwise, blank it out.
Tensor mask = Add(Mul(best_mask, visible), Mul(Sub(1.0f, visible), -1024.0f));
```

Each object now takes a single GPU call per frame. The agent made a series of changes like this, such as folding the attention scale into the weights, and checked each against Hugging Face's SAM 2 before keeping it. That's how we more than tripled the model's frame rate.

We're working on direct [ML Drift](https://developers.googleblog.com/ml-drift-next-gen-gpu-aiml-inference-at-the-edge/) (GPU) backends, NPU-friendly operations, and automatic full-graph quantization. We're also adding Python authoring with `.tflite` export, plus wider dynamic-shape support for varying image sizes, batch sizes, and sequence lengths.

Explore our resources and start your journey with LiteRT:

<sup>The LiteRT Tensor API and its reference implementations are the result of a collaborative effort across the Google AI Edge and LiteRT teams. Special thanks to Byungchul Kim, Changming Sun, Chenchen Tang, Dillon Sharlet, Jingjiang Li, Lu Wang, Lutz Justen, Matthias Grundmann, Misha Gutman, Ronald Wotzlaw, and Sachin Kotwani.</sup>