智东西（公众号：zhidxcom）
编译 | 李水青
编辑 | 云鹏

智东西3月8日报道，今天，谷歌正式发布了MediaPipe LLM Inference API，该API可以让开发人员更便捷地在手机、PC等设备上运行AI大模型。而AI大模型也可以在不同类型的设备上跨设备运行。

谷歌对跨设备堆栈进行了重点优化，包括新的操作、量化、缓存和权重共享等。谷歌称，MediaPipe已经支持了四种模型：Gemma、Phi 2、Falcon和Stable LM，这些模型可以在网页、安卓、iOS设备上运行，谷歌还计划将这一功能扩展到更多平台上。

在安卓上，MediaPipe LLM Inference API仅用于实验和研究，生产应用可以通过安卓AICore在设备上使用Gemini API或Gemini Nano。

这里有一些实时的Gemma 2B通过MediaPipe LLM Inference API运行的动图展示。

▲Gemma-2B在iOS（左）和安卓（右）设备上运行

▲Gemma-2B在设备上的浏览器中运行

GitHub地址：
https://github.com/googlesamples/mediapipe/tree/main/examples/llm_inference
开发者文档地址：
https://developers.google.com/mediapipe/solutions/genai/llm_inference

一、四大模型可用，支持网页、安卓和iOS

MediaPipe LLM Inference API旨在为开发人员简化设备上的大语言模型集成，支持网页、安卓和iOS，最初支持四个公开可用的大模型：Gemma、Phi 2、Falcon和Stable LM，参数规模在1B~3B之间。

与这些架构兼容的任何模型权重都可以与LLM Inference API配合使用。用户可以使用基本模型权重，也可以使用社区微调版或自己的数据微调权重，灵活地在设备上进行原型搭建和测试。

MediaPipe LLM Inference API能够在设备上实现更低的延迟，同时兼容CPU和GPU，支持多个平台。对于高端手机在实际应用中的持续性能表现，安卓AICore可以利用特定的硬件神经加速器进行优化。

▲五大模型在GPU运算时的性能表现

在GPU运算时，Falcon 1B和Phi 2模型采用fp32激活方式，而Gemma和StableLM 3B模型则采用fp16激活方式。因为团队经过质量评估研究发现，后两个模型在精度损失方面表现出更强的稳定性，因此他们为每个模型都选择了能够维持其质量的最低位数激活数据类型。

▲五大模型在GPU运算时的性能表现

二、简单几步实现模型部署，iOS仅能运行Gemma 2B

使用MediaPipe LLM Inference API，用户通过简单几步就能把模型部署到设备上。

以下的代码样本展示了网页（web）SDK，主要步骤如下：

1.选择与谷歌支持的模型架构兼容的模型权重。

2.使用MediaPipe Python包将模型权重转换为TensorFlow Lite Flatbuffer。

3.在应用程序中接入LLM Inference SDK。

4.与应用程序一起托管TensorFlow Lite Flatbuffer。

5.使用LLM Inference API从模型获取文本提示并得到文本响应。

在安卓上，MediaPipe LLM Inference API仅用于实验和研究。生产应用可以通过安卓AICore在设备上使用Gemini API或Gemini Nano。AICore是在安卓14中引入的新的系统级能力，为高端设备提供Gemini驱动的解决方案，包括与最新的ML加速器、用例优化的LoRA适配器和安全过滤器的集成。

在iOS上，由于内存限制，目前只有Gemma 2B（int4）模型可运行，但谷歌正在努力使其他模型也能在iOS平台上启用。

三、权重共享、自定义操作符，六大优化增强模型性能

为了达到更优性能表现，谷歌在MediaPipe、TensorFlow Lite、XNNPack（谷歌CPU神经网络操作库）和GPU加速运行时中进行了无数的优化。

以下是一些带来性能显著提高的一些选定优化：

权重共享：LLM推理通常包含两个阶段：预填充和解码。传统做法是为这两个阶段分别设置独立的推理上下文，各自管理对应的机器学习模型资源。但考虑到大型语言模型对内存的高需求，谷歌引入了一项新特性，允许在不同的推理上下文之间共享权重和键值（KV）缓存。

优化的全连接操作：XNNPack的FULLY_CONNECTED操作在针对大型语言模型推理时，进行了两项关键优化。首先，通过动态范围量化技术，团队成功地将全整数量化在计算和内存方面的优势，与浮点推理在精度方面的优势完美结合。使用int8/int4权重不仅大幅提升了内存吞吐量，还显著提高了性能。其次，团队充分利用了ARM v9 CPU中的I8MM指令，使其可以在一条指令内完成2×8 int8矩阵与8×2 int8矩阵的乘法运算，从而实现了比NEON点积基础实现快一倍的速度。

平衡计算和内存：在对LLM推理进行分析时，团队发现了两个阶段的明显限制：预填充阶段面临计算能力施加的限制，而解码阶段受到内存带宽的限制。因此，每个阶段采用不同的策略对共享的int8/int4权重进行反量化。在预填充阶段，每个卷积操作首先将权重反量化为浮点值，然后进行主要计算，确保计算密集型卷积的最优性能。相反，解码阶段通过将反量化计算添加到主要的数学卷积操作中，最小化了内存带宽。

▲平衡计算和内存过程示意

自定义操作符：对于设备上的GPU加速大模型推理，团队广泛依赖自定义操作来减轻由许多小着色器引起的低效率。这些自定义操作允许进行特殊的操作符融合，并将各种大模型参数（如令牌ID、序列补丁大小、采样参数等）打包到主要用于这些专门操作的专门的自定义张量中。

伪动态性：在注意力块中，团队遇到随着上下文增长而增加的动态操作。由于GPU运行时缺乏对动态操作/张量的支持，团队选择使用预定义的最大缓存大小的固定操作。为了降低计算复杂性，团队引入了一个参数，使得可以跳过某些值计算或处理减少的数据。

优化的KV缓存布局：由于KV缓存中的条目最终作为卷积的权重，代替矩阵乘法，团队将这些存储在为卷积权重定制的专门布局中。这种策略性的调整消除了对额外转换或依赖于未优化布局的必要性，因此有助于更有效和流线型的过程。

结语：谷歌最新API，剑指大模型跨设备运行

早在2019年，谷歌的MediaPipe就已经出现，并开始扩展TensorFlow Lite的能力，起初这些AI工具主要聚焦于小型设备上的模型。

此次谷歌发布的新版本可以让大模型在各个平台上实现完全本地化运行。众所周知，大模型对内存和算力的需求都是很大的，所以谷歌的这项技术颇具变革性。

2024年，谷歌将把MediaPipe LLM Inference API扩展到更多的平台和模型，提供更广泛的转换工具、免费的设备上组件、高级任务等。这也将为其的大模型生态版图扩张提供支持。

来源：谷歌开发者官网