Google 发布 LiteRT Tensor API，用 C++ 手写端侧模型图并跨 CPU、GPU、NPU 运行

OCT. 19, 2026

Ping Yu Staff Software Engineer

Lu Wang Sr. Staff Software Engineer

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, and running it with the LiteRT runtime. 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 made this code-first approach popular on Apple silicon, and 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, 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) 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.

Real-world showcase: Segment Anything on Web

Sorry, your browser doesn't support playback for this video

To show the Tensor API in practice, we built an interactive Segment Anything App, 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), triggered either by simple clicks or natural language prompts powered by Gemma 4 E2B.

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 GitHub repo.

What is the LiteRT Tensor API?

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

• Full control over graph structure: you decide every op in the model graph and every connection between them, from elementwise math to accelerator-optimized LLM composite ops, such as scaled dot-product attention, RMSNorm, and KV-cache updates.
• Direct weight ingestion: load weights straight from model checkpoints, such as Hugging Face .safetensors files, without conversion scripts needed in between.
• Tensor-centric buffer interop: pre-processing, model partitions, and post-processing pass tensors to each other through shared accelerator buffers, so intermediate data stays on the accelerator with zero copies.
• Building blocks for AI agents: models are plain C++ built from well-defined ops, so coding agents can construct a model, check its results against a reference, and tune its performance by benchmarking variants on each device.

Cross-platform performance

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 includes the benchmark commands so you can reproduce the results.

Note: Times cover the image encoder and mask decoder; pre- and post-processing add ~0.34 ms. Output masks match Meta's PyTorch reference pixel for pixel.

Get started with Tensor API

Authoring model in native C++

With the Tensor API, each layer is a C++ function over Tensor 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"));
}

C

Copied

In the SAM 2 sample, BuildEncoder and BuildDecoder compose layers like this into the encoder and decoder graphs, and ModelFactory 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 (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();

C

Copied

SAM 2 takes the .tflite path to reach WebGPU in the browser, and the rest of this post follows it. The Gemma 4 example takes the direct path for prefill and decode on CPU, the lightest way to run Tensor API code. The Tensor API README 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.

Expressing the full ML computation graph

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.

Zero copies: Keep the whole pipeline on the 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, all four steps are Tensor API graphs on the GPU. Pre-processing, for example, is a small graph built with CreateLambdaRunner. 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"));

C

Copied

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.

Full control: Tune the graph for each backend

Because the graph is plain C++, you or your coding agent can rewrite it for each backend. Our agent, Antigravity, 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:

// 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));

C

Copied

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.

What's next

We're working on direct ML Drift (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:

• Using a coding agent? Point it to the Tensor API README and the SAM 2 sample; together they cover the API and a buildable end-to-end example.
• Documentation: LiteRT Tensor API guide (getting started, build instructions and examples) and Tensor API reference.
• Reference Implementation: SAM 2 authored in C++ on the LiteRT Tensor API, Live Web Demo, and Web App Source.
• LLM examples: Gemma 4 running directly on XNNPACK (CPU) and Gemma 3 running through LiteRT (CPU and GPU), both authored with the Tensor API.
• Model Definition: sam2_graph.cc - the Hiera encoder and mask decoder graphs.

Acknowledgements

^{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.}

Previous

Next