Google 发布 LiteRT Tensor API，支持用 C++ 手写端侧模型计算图
Native Model Authoring for Edge with LiteRT Tensor API
Google 发布 LiteRT Tensor API，这是 LiteRT 栈中的轻量建模层，开发者可用 C++ 逐算子构建计算图，同一份代码经 LiteRT 跑在 CPU、GPU（ML Drift）和 NPU 上，或经 XNNPACK 直接在 CPU 运行、无需模型文件。
Google 官方给出 LiteRT Tensor API 的设计与代码示例，读者可据此判断手写端侧推理图与转换路径的取舍。
2026年10月19日
Ping Yu 资深软件工程师
Lu Wang 高级主任软件工程师
在设备端部署 AI 意味着要在每个目标设备上满足严苛的延迟和内存预算。在生产环境中部署模型最常见的路径是转换训练好的模型，例如使用 LiteRT Torch 转换 PyTorch 模型，然后用 LiteRT 运行时运行。这是通往生产环境最快捷、最可靠的途径，也是我们为大多数项目推荐的路径。不过，有时开发者希望完全掌控网络的表示和优化方式，以支持新的架构，或尝试新的 lowering 逻辑，而不必等待上游框架更新。
为了实现终极控制，开发者会用原生 C++ 或着色器语言手动编写计算。 Apple 的 MLX 让这种代码优先的方法在 Apple silicon 上流行起来，而 llama.cpp 则展示了定制 C++ 实现能把 LLM 推理推进到多远。然而，这种级别的控制代价不菲：实现与特定后端（如 Metal、OpenCL 或 Vulkan）紧密耦合，每换一个新硬件目标就必须重写。此外，管理 CPU、GPU 和 NPU 之间的零拷贝缓冲区互操作与同步，同时还要把预处理、推理和后处理粘合在一起，需要大量手动工程工作。
为了解决上述挑战，我们很高兴地宣布推出全新的 LiteRT Tensor API，这是 LiteRT 技术栈中的轻量级编写层，旨在改进手工构建模型的方式。它结合了原生编写的灵活性与跨平台可移植性。你用 C++ 逐个算子构建计算图，精确选择运行哪些算子以及它们如何连接。同一个图可以通过 LiteRT 在 CPU、GPU（由开源 ML Drift 驱动）和 NPU 上运行，或者无需模型文件和运行时，直接通过 XNNPACK 在 CPU 上运行，因此不必为每个新硬件目标重写。模型推理和前后处理可以用同一套 API 组装，并通过共享加速器缓冲区串联起来，使数据无需往返 CPU。请从转换路径开始，这仍然是大多数模型的最佳默认选择。当你需要对图进行底层控制，或希望整个流水线端到端在加速器上运行时，再使用 Tensor API。
LiteRT Tensor API 还是AI 编码代理的理想编程基础，提供结构化原语来自动生成、重新连接和调优模型图。在本文的剩余部分，我们将深入探讨其设计，并通过实际代码示例演示如何从零开始编写模型。
真实案例展示：Web 上的 Segment Anything
抱歉，你的浏览器不支持此视频的播放
为了展示 Tensor API 的实际应用，我们构建了一个交互式 Segment Anything 应用，完全在浏览器中通过 WebGPU 在设备端运行。该应用通过 Meta 的 Segment Anything Model 2 (SAM 2) 实现实时多对象分割与跟踪，可由简单点击或由 Gemma 4 E2B 驱动的自然语言提示触发。
SAM 2 模型完全通过 LiteRT Tensor API 用 C++ 编写。一个编码智能体编写并对其进行了演进，使其在 WebGPU 上以 512x512 分辨率从不到 10 帧/秒提升到超过 30 帧/秒。完整实现和 Web 应用已在 LiteRT Samples GitHub 仓库中开源。
什么是 LiteRT Tensor API？
LiteRT Tensor API 提供四项关键功能，用于编写和优化端侧模型：
- 完全掌控图结构：模型图中的每个算子以及它们之间的每条连接都由你决定，从逐元素数学运算到针对加速器优化的 LLM 复合算子，例如缩放点积注意力、RMSNorm 和 KV-cache 更新。
- 直接读取权重：直接从模型检查点加载权重，例如 Hugging Face 的 .safetensors 文件，无需中间的转换脚本。
- 以张量为中心的缓冲区互操作：预处理、模型分区和后处理通过共享的加速器缓冲区相互传递张量，因此中间数据留在加速器上，零拷贝。
- 面向 AI 智能体的构建模块：模型是由定义良好的算子构建的纯 C++，因此编码智能体可以构建模型、对照参考结果检查其输出，并通过在每个设备上对变体进行基准测试来调优其性能。
跨平台性能
我们在三台设备的 GPU 上运行了 SAM 2 模型的同一份 C++ 代码：一台 macOS 笔记本电脑（Metal）、一部 iOS 手机（Metal）和一部 Android 手机（OpenCL）。在我们于 MacBook Pro 上的测试中，它比 Apple MLX 快约 1.5 倍（fp16，图像编码器加掩码解码器）。SAM 2.1 示例包含基准测试命令，你可以复现这些结果。
注意：时间涵盖图像编码器和掩码解码器；预处理和后处理增加约 0.34 ms。输出掩码与 Meta 的 PyTorch 参考结果逐像素一致。
开始使用 Tensor API
用原生 C++ 编写模型
使用 Tensor API，每一层都是作用于 Tensor 值的 C++ 函数，由高层算子和直接从检查点加载的权重构建。例如，下面是一个多头自注意力层：
// 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
已复制
在 SAM 2 示例中，BuildEncoder 和 BuildDecoder 将这样的层组合成编码器和解码器图，而 ModelFactory 将两者导出为一个 sam2_image.tflite 文件的两个签名。
.tflite 文件是运行图的两种方式之一。后端是 Tensor: TfLiteMixinTag 上的模板参数，导出 .tflite，由 LiteRT 在 CPU、ML Drift（GPU）或 NPU 上运行，而 XnnpackMixinTag 通过 XNNPACK 直接在 CPU 上运行相同代码，无需模型文件。为张量同时打上两者标记即可保留两种选项：
// 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
已复制
SAM 2 采用 .tflite 路径以在浏览器中达到 WebGPU，本文其余部分也遵循这一路径。Gemma 4 示例采用直接路径在 CPU 上进行预填充和解码，这是运行 Tensor API 代码最轻量的方式。Tensor API README 对两者都有介绍。
一种自然的工作流会同时使用两者：在直接的 XNNPACK 路径上迭代，直接从 .safetensors 检查点加载权重，无需重建模型文件；然后在准备发布时导出 .tflite，在移动端通过 LiteRT 或在 Web 上通过 LiteRT.js 运行。
表达完整的机器学习计算图
一个真实的应用需要的不仅仅是神经网络本身。在演示中，将摄像头帧和一次点击转换为掩码需要四个步骤：预处理将帧转换为模型期望的 512x512 图像；图像编码器将该图像转换为特征；掩码解码器将这些特征与点击结合以预测掩码；后处理将掩码缩放回帧的尺寸。要快速运行整个流水线需要两件事：在步骤之间将数据保留在加速器（如 GPU）上，以及控制每个步骤在每个加速器上的运行方式。
零拷贝：将整个流水线保留在加速器上
在典型的应用中，只有模型在 GPU 上运行；CPU 调整每个摄像头帧的尺寸，将编码器的输出传递给解码器，并对掩码进行后处理。每次往返都会复制数据并导致 GPU 停顿。
在 litert-samples/sam2 中，所有四个步骤都是 GPU 上的 Tensor API 图。例如，预处理是用 CreateLambdaRunner 构建的一个小图。这些图在启动时连接一次以共享缓冲区，然后针对每一帧和点击运行：
// 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
已复制
步骤之间没有数据返回 CPU：所有 522 个编码器操作和 269 个解码器操作都在 GPU 上运行，预处理和后处理每帧仅增加约 0.34 毫秒。
完全控制：为每个后端调优图
由于图是纯 C++ 的，你或你的编码代理可以为每个后端重写它。我们的代理 Antigravity 将 SAM 2 的逐帧决策（例如对象是否仍然可见）从 CPU 代码移入图中。像 score > 0 这样的比较在 WebGPU 上仍会回退到 CPU，因为它们输出 true/false 值，因此代理将它们重写为 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
已复制
现在每个对象每帧只需一次 GPU 调用。代理进行了一系列类似的更改，例如将注意力缩放折叠到权重中，并在保留每项更改之前对照 Hugging Face 的 SAM 2 进行检查。这就是我们将模型的帧率提高三倍多的方式。
下一步
我们正在开发直接的 ML Drift（GPU）后端、对 NPU 友好的操作以及自动全图量化。我们还在添加使用 .tflite 导出的 Python 编写功能，以及对可变图像尺寸、批大小和序列长度的更广泛的动态形状支持。
探索我们的资源，开始你的 LiteRT 之旅：
- 使用编码代理？ 将其指向 Tensor API README 和 SAM 2 示例；它们共同涵盖了 API 和一个可构建的端到端示例。
- 文档：LiteRT Tensor API 指南（入门、构建说明和示例）以及 Tensor API 参考。
- 参考 实现：在 LiteRT Tensor API 上用 C++ 编写的 SAM 2、在线 Web 演示 和 Web 应用源码。
- LLM 示例：Gemma 4 直接在 XNNPACK（CPU）上运行，以及 Gemma 3 通过 LiteRT（CPU 和 GPU）运行，两者均使用 Tensor API 编写。
- 模型 定义：sam2_graph.cc - Hiera 编码器和掩码解码器图。
致谢
LiteRT Tensor API 及其参考实现是 Google AI Edge 和 LiteRT 团队通力合作的成果。特别感谢 Byungchul Kim、Changming Sun、Chenchen Tang、Dillon Sharlet、Jingjiang Li、Lu Wang、Lutz Justen、Matthias Grundmann、Misha Gutman、Ronald Wotzlaw 和 Sachin Kotwani。
上一篇
下一篇
来源：Google Developers Blog · developers.googleblog.com