自 2017 年发布以来，TensorFlow Lite 一直是实现设备端机器学习的强大工具，而 MediaPipe 在 2019 年通过支持完整的 ML 管道进一步增强了这一能力。尽管这些工具最初侧重于较小的设备端模型，但今天，实验性的 MediaPipe LLM 推理 API 的推出标志着机器学习领域迎来了重大变革。
此次发布的新版本使得大型语言模型 (LLM) 能够在各平台上完全实现设备端运行。考虑到 LLM 的内存及计算需求是传统设备端模型的百倍之多，这一新功能具有重大变革意义。实现 LLM 全设备运行的关键在于对整个设备端技术栈进行了深度优化，其中涉及新的运算符、量化、缓存以及权重共享等方面。
实验性的跨平台 MediaPipe LLM 推理 API 旨在简化 Web 开发者在设备端集成 LLM 的过程。该 API 支持 Web、Android 及 iOS 平台，并初步兼容四个公开可用的 LLM 模型：Gemma、Phi 2、Falcon 以及 Stable LM。借助该 API，研究人员与开发者能灵活地在设备端构建原型并测试那些公开可用的热门 LLM。
在 Android 平台上，MediaPipe LLM 推理 API 仅限于实验和研究用途。具有 LLM 的正式版应用可使用 Gemini API 或通过 Android AICore 在设备端使用 Gemini Nano。AICore 是 Android 14 引入的一项系统级功能，旨在为高端设备提供由 Gemini 驱动的解决方案。这些解决方案包括与最新 ML 加速器的集成、针对特定用例优化的 LoRA 适配器以及安全过滤条件等。若希望在您的应用中开始使用 Gemini Nano 进行设备端部署，请申请加入抢先体验计划。
从今天起，您可以通过使用网页演示或构建示例演示应用来试用 MediaPipe LLM 推理 API。您可以通过我们的 Web、Android 或 iOS SDK 将该 API 融入您的项目进行试验和集成。
借助 LLM 推理 API，只需几个步骤即可将 LLM 部署至设备端。虽然 SDK 和原生 API 会根据平台（Web、iOS 和 Android）的不同而有所差异，但部署步骤在这些平台上是通用的。以下代码示例展示了 Web SDK 的用法。
2. 使用 MediaPipe Python 软件包将模型权重转换为 TensorFlow Lite Flatbuffer
from mediapipe.tasks.python.genai import converter
config = converter.ConversionConfig(...)
converter.convert_checkpoint(config)
3. 在您的应用中包含 LLM 推理 SDK
import { FilesetResolver, LlmInference } from "https://cdn.jsdelivr.net/npm/@mediapipe/tasks-genai”
4. 将 TensorFlow Lite Flatbuffer 与您的应用一起托管。
5.使用 LLM 推理 API 获取文本提示并从模型中获取文本响应
const fileset = await FilesetResolver.forGenAiTasks("https://cdn.jsdelivr.net/npm/@mediapipe/tasks-genai/wasm");
const llmInference = await LlmInference.createFromModelPath(fileset, “model.bin”);
const responseText = await llmInference.generateResponse("Hello, nice to meet you");
document.getElementById('output').textContent = responseText;
我们的初始版本支持以下四种模型架构。任何与这些架构兼容的模型权重都将与 LLM 推理 API 兼容。您可以使用基础模型权重，也可以使用社区提供的微调版本权重，或是使用您自己的数据对权重进行微调。
经过大量优化工作（其中部分内容将在下文详述），MediaPipe LLM 推理 API 能够在设备端实现业界顶尖的低延迟性能。该性能重点关注 CPU 和 GPU，以确保 API 能够在多个平台上运作。在部分高端手机的生产环境中，Android AICore 能够利用特定于硬件的神经加速器以保证持续稳定的性能。
衡量 LLM 延迟时，需要关注以下几个术语和指标。其中，“首个令牌时间”(Time to First Token) 和“解码速度”(Decode Speed) 最为重要，它们分别衡量您开始获得响应的速度以及响应开始后的生成速度。
“预填充速度”(Prefill Speed) 和“解码速度”(Decode Speed) 受模型、硬件以及最大令牌数的影响，并且会随设备当前负载的变化而变化。
以下速度值是在高端设备上测量而得的，使用了 1280 的最大令牌数、1024 个令牌的输入提示长度，以及 int8 权重量化。Gemma 2B (int4) 模型（可在此 Kaggle 页面 上找到）是个例外情况，它采用了混合 4/8 位权重量化。
为了实现上述性能表现，我们在 MediaPipe、TensorFlow Lite、XNNPack（我们的 CPU 神经网络运算符库）以及 GPU 加速运行时等多个方面进行了大量优化工作。以下是其中一些带来了显著性能提升的精选优化措施。
权重共享：LLM 推理过程包含 2 个阶段：预填充阶段和解码阶段。传统而言，此设置需要 2 个独立的推理上下文，每个上下文各自管理其对应 ML 模型的资源。鉴于 LLM 对内存的需求，我们添加了一项功能，允许模型在推理上下文之间共享权重和 KV 缓存。虽然共享权重看似简单，但在计算密集型和内存密集型运算之间进行共享会对性能产生重大影响。在典型的 ML 推理场景中，权重不会与其他运算符共享，而是为每个全连接运算符单独精心配置，以确保最佳性能。与另一个运算符共享权重可能会导致各运算符无法充分优化。为解决此问题，开发者需要编写新的内核实现，使其即使在次优权重下也能高效运行。
优化的全连接运算：针对 LLM 推理，我们对 XNNPack 的 FULLY_CONNECTED 运算进行了两项重大优化。首先，动态范围量化无缝融合了全整数量化在计算和内存方面的优势以及浮点推理的精度优点。使用 int8/int4 权重不仅能提高内存吞吐量，还能实现卓越性能，特别是对于 4 位权重，仅需一条额外指令便能在寄存器内实现高效解码。其次，我们积极利用 ARM v9 CPU 中的 I8MM 指令，在单个指令中将 2x8 int8 矩阵乘以 8x2 int8 矩阵，从而达到比基于 NEON 点积实现方法快两倍的速度。
平衡计算与内存资源：在对 LLM 推理进行分析后，我们发现预填充阶段和解码阶段分别面临不同的限制：预填充阶段受限于计算能力，而解码阶段则受制于内存带宽。因此，我们这两个阶段采用了不同的策略来对共享的 int8/int4 权重进行反量化。在预填充阶段，每个卷积运算符会在进行主要计算之前先将权重反量化为浮点值，以确保在计算密集型卷积运算中达到最佳性能。相反，在解码阶段，我们将反量化计算与主要的数学卷积运算相结合，以最小化内存带宽。
自定义运算符：为了在设备端利用 GPU 加速 LLM 推理，我们在很大程度上依赖自定义运算来缓解由于众多小型着色器引起的效率低下问题。这些自定义运算符支持将特殊运算符融合以及将各种 LLM 参数（如令牌 ID、序列补丁大小、采样参数等）打包到主要用于这些专业运算的专用自定义张量中。
伪动态：在注意力块中，随着上下文增长，我们会遇到随时间递增的动态运算。由于我们的 GPU 运行时不支持动态运算/张量，我们选择使用具有预定义最大缓存大小的固定运算。为了降低计算复杂性，我们引入了一个参数，以跳过某些值的计算或处理缩减数据。
优化的 KV 缓存布局：由于 KV 缓存中的条目最终作为卷积的权重使用，替代了矩阵乘法，我们将其存储在一个专为卷积权重设计的特殊布局中。这一策略调整消除了对额外转换或依赖于未优化布局的需求，从而有助于使流程变得更高效、更流畅。
我们对 MediaPipe LLM 推理 API 实验性版本的各项优化及其展现出的性能感到非常激动。但这只是项目发展的一个起点，在整个 2024 年里，我们将继续扩展支持更多平台与模型，提供更丰富的转换工具、配套的设备端组件，并支持更高级的任务。敬请期待更多成果。
您可以在 GitHub 上查看官方示例，其中涵盖了前文所介绍的所有内容。如需获取更多细节，欢迎阅读我们的官方文档。请密切关注 Google for Developers YouTube 频道，以获取更新内容和教程。
我们要感谢所有为本工作做出贡献的团队成员：T.J. Alumbaugh、Alek Andreev、Frank Ban、Jeanine Banks、Frank Barchard、Pulkit Bhuwalka、Buck Bourdon、Maxime Brénon、Chuo-Ling Chang、Lin Chen、Linkun Chen、Yu-hui Chen、Nikolai Chinaev、Clark Duvall、Rosário Fernandes、Mig Gerard、Matthias Grundmann、Ayush Gupta、Mohammadreza Heydary、Ekaterina Ignasheva、Ram Iyengar、Grant Jensen、Alex Kanaukou、Prianka Liz Kariat、Alan Kelly、Kathleen Kenealy、Ho Ko、Sachin Kotwani、Andrei Kulik、Yi-Chun Kuo、Khanh LeViet、Yang Lu、Lalit Singh Manral、Tyler Mullen、Karthik Raveendran、Raman Sarokin、Sebastian Schmidt、Kris Tonthat、Lu Wang、Zoe Wang、Tris Warkentin、Geng Yan、Tenghui Zhu，以及 Gemma 团队。