Native Model Authoring for Edge with LiteRT Tensor API

Google released the LiteRT Tensor API, a C++ authoring layer in the LiteRT stack that lets developers build computation graphs op by op and run the same graph on CPU, GPU via open-source ML Drift, and NPU, or directly on CPU through XNNPACK without a model file or runtime. Google demonstrated the API with an on-device Segment Anything Model 2 (SAM 2) app running in the browser on WebGPU, where a coding agent evolved the C++ implementation from under 10 to over 30 frames per second at 512x512, with code published in the LiteRT Samples GitHub repo. Google said the conversion path, such as converting a PyTorch model with LiteRT Torch, remains the recommended default for most projects, and positions the Tensor API as a programmatic foundation for AI coding agents to generate, rewire, and tune model graphs.

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