端侧AI架构解析：模型、框架与硬件

半行代码
94 阅读36分钟

最近了解了一些端侧 AI 推理的技术方案，在此做一些分享。

随着 AI技术 和大模型的快速发展，AI 已经成为各个业务场景里探索创新的重要方向。随着端侧算力的不断提升和端侧模型性能的不断增强，端侧 AI 部署及推理成为了新的热门话题之一。本篇文章对市面上常见的端侧推理引擎、端侧大模型部署方案做一个介绍。

端智能的优势

优势

端智能在端侧本地进行推理和决策，不依赖云端服务器，相比传统对接云端智能，存在不少优势：

  1. 低延迟
  • 模型运行在本地，不需要网络传输，响应速度更快
  • 一些实时性比较高的场景非常适合，例如智能辅助驾驶
  1. 隐私保护
  • 用户本地数据不需要上传云端，减少用户隐私泄漏的风险
  1. 离线可用
  • 在弱网甚至无网状态下也能正常运行
  • 大大提升用户体验和系统可靠性
  1. 降低云端成本
  • 不需要在云端进行推理，大大节省云端的计算资源和计算成本
  • 减少数据上传云端，节省网络带宽

算力

端侧推理还有一个绕不开的话题就是算力。我整理了常见类型的端侧设备的算力，用 TOPS(每秒万亿次操作次数) 表示

  • 智能手机

智能手机的主流算力范围大约为10～50 TOPS，典型代表为苹果A系列、高通骁龙、联发科天玑、华为麒麟芯片

  • 智能汽车

智能汽车一般分为座舱芯片和智驾系统芯片，座舱一般为高通SA系列，例如8155P、8295P、8775P，算力分别为8 TOPS、30 TOPS和70 TOPS。智驾芯片之前很多车型使用英伟达 NVIDIA Orin芯片，单芯片算力一般在 250 TOPS，近年来国产芯片大量上车，例如蔚来神玑芯片(单颗1000 TOPS)、小鹏图灵芯片(单颗750 TOPS)。

  • 可穿戴设备，例如智能眼镜、智能手表

常见旗舰款智能眼镜常见芯片为高通AR1+ Gen1，算力约1 TOPS，旗舰款智能手表芯片常见芯片有苹果、华为麒麟A2、高通Snapdragon W5，算力基本不到1 TOPS

  • 智能家居，IOT设备

晶晨S905X5、全志V853等，算力基本只有几 TOPS

目前来看，各种各样的端设备的芯片算力其实都不低。在一些场景上算力还相当的高。端智能在当下和未来都会有非常好的落地前景。

模型量化和蒸馏

机器学习模型需要部署在端侧，一般需要对模型的大小进行优化。常见的优化手段包括模型量化、模型知识蒸馏。在主流机器学习框架，包括 PyTorch、TensorFlow、HuggingFace 等，对量化和蒸馏都有较好的支持，本文目的在于描述清楚这两个优化操作的基础知识。

模型量化

模型量化是模型入端之前的常见优化之一，目的是将模型权重和激活值从高bit宽度转换成低bit宽度表示。从而显著降低i计算成本和内存开销。例如将FP16的浮点张量转成INT8来表示。

量化方法分类

按量化时机分类

模型量化方法根据训练后和训练中划分为训练后量化和量化感知训练。

  • 训练后量化：针对已完成训练的模型进行量化，不需要自备训练数据集，只需要少量的校准数据集，避免高昂的计算成本。但是模型每一层量化是单独进行的，训练后量化有可能导致精度损失变大。
  • 感知量化训练：在训练中模拟量化操作，模型会学会在量化条件下依然正常工作。优点是精度高，推理结果准确，缺点则是计算成本高。

一般来说入端之前模型会经过训练后量化得到量化后的新模型。量化后模型会减少内存占用，提升推理速度。

按函数形状分类

模型量化方法还可以按照量化函数形状分为线形量化和非线形量化。

  • 线形量化：采用相同的量化间隔对输入做量化。比如FP16的精度范围通过平移、缩放线形的映射到INT8的[-128，127] 范围。
  • 非线形量化：将浮点数数值分布范围非均匀划分，通常小值划分精细，大值划分粗糙，使用查表、对数、分段函数等方式进行映射

两者的目的都是为了不同数据分布的情况下，尽量减少精度损失。

按零点分类

按零点分类则把模型量化方法分为对称量化和非对称量化。介绍此分类方法前先介绍何为零点，零点是非对称量化里的一个关键参数，它的核心作用是将浮点数里的“0”映射到整数量化范围内的某个整数，来支持非对称数据的分布。

  • 对称量化：零点为真正的0，即浮点0映射后为整数0，输入数据是分布在0的两侧且均匀分布。输入数据直接映射为量化后数据。这里映射可以按线形映射也可以是非线形映射。
  • 非对称量化：一些激活值计算结果的范围基本都在0以上，没有负数，如果直接映射到整数值例如INT8，那么就会导致[-128,0]这个范围没有数据，从而浪费了INT8一半的精度，使得精度下降。引入零点概念后，零点作为映射运算的参数（类似偏移量），把数据重新映射到[-128-127]的范围里。

大模型量化方法

而在大模型的量化中，会根据 prefill 和 decode 2个阶段采取不同的量化手段。

prefill

prefill 阶段会直接计算所有 token，而不是一个一个逐步计算，主要的数学计算为通用矩阵乘法。计算的延迟主要收计算精度影响，所以 prefill 阶段会将权重和激活值一起量化：

decode

decode 发生在大模型的自回归阶段，每个步骤只处理一个token，所以主要进行的是矩阵向量乘法运算。计算的延迟主要受加载大模型权重，所以 decode 阶段会量化权重，一般不量化激活值。

模型知识蒸馏

另一个常见的模型精简方法叫做知识蒸馏。核心思想是将一个大型模型的知识迁移到一个比较小的模型。这个大型模型叫做教师模型，较小的模型叫做学生模型。根据是否可以访问大模型内部结构，知识蒸馏分为白盒蒸馏和黑盒蒸馏。

白盒知识蒸馏

白盒知识蒸馏中，学生模型可以完全访问教师模型的内部结构和输出，包括教师模型的输出、中间层的特征表示、梯度信息、logtis等。学生可以模仿教师等中间表示，而不仅仅是最终预测，可以学到更细粒度的知识。

黑盒知识蒸馏

黑盒知识蒸馏中，学生模型只能访问教师模型的输入和输出。黑盒蒸馏的特点是保护教师模型的隐私和知识产权。

知识蒸馏的流程如下，中间特征损失只有白盒知识蒸馏才有：

端侧模型推理

ONNX

ONNX和onnxruntime

ONNX(Open Neural Network Exchange，开放式神经网络交换格式)，是一种开放、通用的深度学习模型表示标准。由微软等公司开发，目前由ONNX开源社区维护。ONNX 允许在 Pytorch、TensorFlow等主流训练框架中训练的模型统一导出为 .onnx 格式的模型文件。

因为各家对于 onnx 格式都有支持，所以 onnx 事实上已经成为了不同模型训练推理框架互相转换的一个中间格式。

以下是一个 BERT 模型从 Pytorch 转成 .onnx 格式的 Python 代码：

import torch
from transformers import AutoModel, AutoTokenizer, AutoConfig, BertForMaskedLM

def main():
  tokenizer = AutoTokenizer.from_pretrained('bert-base-chinese')
  model = BertForMaskedLM.from_pretrained('bert-base-chinese')
  model.eval()
  text='北京是[MASK]国的首都'
  inputs = tokenizer(text, return_tensors='pt')
  input_names = ['input_ids','attention_mask', 'token_type_ids']
  output_names = ['logits']
  //导出.onnx
  torch.onnx.export(
    model,
    args=(inputs["input_ids"], inputs["attention_mask"], inputs["token_type_ids"]),
    f='results/onnx/model.onnx',
    input_names=input_names,
    output_names=output_names,
    do_constant_folding=True,
    opset_version=14,
    dynamic_axes={
      "input_ids": {0: "batch", 1: "sequence"},
      "attention_mask": {0: "batch", 1: "sequence"},
      "token_type_ids": {0: "batch", 1: "sequence"},
      "logits": {0: "batch", 1: "sequence"},
    },
    training=torch.onnx.TrainingMode.EVAL,
    verbose=False
  )

.onnx文件内部结构基于 Protocol Buffers 定义，核心是一个 GraphProto 对象。包括了

  • inputs：输入张量列表
  • outputs：输出张量列表
  • node：算子节点列表
  • initializer：常量权重

通过netron.app可以看到 .onnx 格式的图结构，可以看到一个BERT模型的图结构就已经非常庞大了：

ONNX同时也支持设置动态 shape，否则 shape 会被固定，推理的时候传入不一致的 shape 就会报错。下图是 BERT 的动态 shape 设置，可以看到张量的 shape 显示的都是变量名：

ONNX 提供了端侧推理引擎 ONNX Runtime，支持多个平台和多个开发语言，在 Android 上也有 aar 提供，支持我们调用 Java API 直接进行推理。还是以 Android 工程里执行 BERT 的推理为例：

private val session: OrtSession = OrtEnvironment.getEnvironment().createSession(modelFile.absolutePath, OrtSession.SessionOptions())

fun forward(text: String) {
  val tokens: List<String> = tokenizer.tokenize(text)
  val inputIds : LongArray = tokens.map { tokenizer.tokenToId(it).toLong() }.toLongArray()
  val seqLen = inputIds.size
  val attentionMask = LongArray(seqLen){1}
  val tokenTypeIds = LongArray(seqLen) {0}
  val environment = OrtEnvironment.getEnvironment()

  val inputIdsTensor = OnnxTensor.createTensor(
    environment, arrayOf(inputIds)
  )
  val attentionMaskTensor = OnnxTensor.createTensor(
    environment, arrayOf(attentionMask)
  )
  val tokenTypeIdsTensor = OnnxTensor.createTensor(
    environment, arrayOf(tokenTypeIds)
  )

  val ortInputs = mapOf(
    "input_ids" to inputIdsTensor,
    "attention_mask" to attentionMaskTensor,
    "token_type_ids" to tokenTypeIdsTensor
  )

  // 寻找[MASK] 位置
  val maskIndex = inputIds.indexOf(103)
  // 推理
  val outputs = session.run(ortInputs)
  val logitsTensor: OnnxTensor? = (outputs.get("logits") as? Optional<OnnxTensor>)?.get()
  logitsTensor?.let {logitsTensor->
    val logits = logitsTensor.value as Array<Array<FloatArray>>
    val maskLogits = logits[0][maskIndex]
    val probs = softmax(maskLogits)
    // 取出top5 [MASK] 位置概率
    val topK = getTopK(probs, 5)
    val results = topK.map { (id, prob) ->
      tokenizer.idToToken(id) to prob // 词表映射
    }
  }
}

onnxruntime内存管理与架构设计

onnxruntime通过多层内存分配来管理内存：

  • 上层叫做 Arena 分配器，会预分配大内存块按需给张量。推理过程中不会频繁分配内存，内存可以复用。
  • 底层叫做设备分配器，权重内存、Arena预分配的大内存块都通过设备内存分配器分配，抽象这一层分配器是为了用一个统一的接口来处理不同设备（CPU、GPU）上的内存分配和释放操作。

onnxruntime的内存管理策略，减少了内存的波动。设计上可配置性，可扩展性比较强。缺点就是初始化有内存规划的开销，Arena也有可能会存在内存浪费。

onnnxruntime 支持通过 Execution Provider 机制提供各个平台的推理实现，Execution Provider 把 onnx 模型映射到特定硬件上的后端插件。例如 NVIDIA 的 CUDA、通过 NNAPI 跑在 Android、通过 CoreML 跑在 iOS、通过 QNN 跑在高通芯片。

但是某些硬件平台，例如 NPU，首次运行模型的时候需要对模型进行编译转换，这个过程可能非常耗时。为了解决这个问题，onnxruntime 引入了 EP Context 机制。将一次性的、耗时的模型编译结果进行缓存，下次推理的时候直接加载缓存，跳过重复编译。

Google AI Edge - LiteRT

Google推出了边缘计算的推理框架LiteRT，LiteRT支持高性能ML，并且通过LiteRT-LM支持生成式 AI，也支持使用 MediaPipe Task 使用现成的机器学习任务。LiteRT 是过去的 TensorFlow Lite 框架的继承者，支持部署 .tflite 模型

LiteRT 经过获取/训练/转换模型、量化/蒸馏模型、模型加载/推理、选择后端执行环境几个步骤。

推理流程

在 Android 里， 可以使用 Kotlin 代码编写推理过程， LiteRT 进行推理的 Kotlin 代码如下：

val compiledModel = CompiledModel.create("/path/to/mymodel.tflite",CompiledModel.Options(Accelerator.CPU))
val inputBuffers = compiledModel.createInputBuffers()
val outputBuffers = compiledModel.createOutputBuffers()
inputBuffers.get(0).writeFloat(input0)
inputBuffers.get(1).writeFloat(input1)
// 推理并读取输出
compiledModel.run(inputBuffers, outputBuffers)
val output = outputBuffers.get(0).readFloat()

这里 inputBuffers 是一个 List，用来存放输入向量。对应上节的 BERT 模型，就是把 inputIds、inputMask、tokenTypeIds 通过 TensorBuffer#writeFloat 写入，伪代码如下：

// 输入转成 FloatArray
inputBuffers[0].writeFloat(intArrayToFloatArray(inputIds))
inputBuffers[1].writeFloat(intArrayToFloatArray(inputMask))
inputBuffers[2].writeFloat(intArrayToFloatArray(tokenTypeIds))

硬件加速与内存管理

LiteRT 支持在 GPU 和 NPU 上进行硬件加速。

  • 借助 LiteRT 的 GPU 加速功能，您可以创建 GPU 友好的输入和输出缓冲区，在 GPU 内存中实现数据零拷贝，并异步执行任务以最大限度地提高并行性。
  • LiteRT 提供了一个统一的接口来 NPU，不需要用户自己逐步对接芯片厂商。使用 LiteRT 进行 NPU 加速可提升实时推理和大型模型推理的性能，并通过使用零拷贝硬件缓冲区来最大限度地减少内存拷贝。LiteRT 的 NPU 加速支持高通 QNN 和 联发科的 NeuroPilot。LiteRT NPU 支持 AOT 和 JIT 两个模式，AOT 模式会把合适的目标 SoC 已知的大型复杂模型提前编译，适合比较复杂的大模型，JIT 模式在模型初始化的时候进行编译，首次运行成本稍高，适合小模型。

LiteRT 通过 Delegate 机制实现硬件加速能力。Delegate 设计成一种可插拔接口，每一类 Delegate 实现都对应了特定的优化后端。在模型加载阶段，LiteRT 会把计算图中可加速的子图委托给对应的 Delegate，Delegate 会把子图编译为硬件指令，并且在推理的时候调用对应的实现。不支持的算子就自动回退到 CPU 进行计算。

LiteRT 设计了多种内存分配策略来优化内存分配

typedef enum TfLiteAllocationType {
    kTfLiteMemNone = 0,
    kTfLiteMmapRo,
    kTfLiteArenaRw,
    kTfLiteArenaRwPersistent,
    kTfLiteDynamic,
    kTfLitePersistentRo,
    kTfLiteCustom,
    kTfLiteVariantObject,
    kTfLiteNonCpu,
} TfLiteAllocationType
类型含义用途
kTfLiteMmapRo只读内存映射加载模型权重，节省内存
kTfLiteArenaRw临时读写的内存池中间激活值
kTfLiteArenaRwPersistent持久化读写内存池推理过程中需要保存的状态
kTfLiteDynamic动态分配内存动态shape张量
kTfLitePersistentRo准备阶段预计算张量内存模型准备阶段计算的张量
kTfLiteCustom自定义内存外部传入缓冲区，例如摄像头帧
kTfLiteVariantObjectC++类型擦出对象存储任意C++对象
kTfLiteNonCpu非CPU内存GPU/NPU的零拷贝张量

硬件加速

GPU加速

当推理计算在 CPU 上执行的时候，CPU 通常以逐层顺序调度的方式执行推理图，每个算子由高度优化的 SIMD 指令或者多线程进行一定程度的并行加速，但是并行速度有限。CPU 的核心数量一般也比较少，只有4-8核，遇到计算密集型运算任务比较吃力。而 GPU 拥有非常多数量的并行计算单元，非常适合这类计算密集型的推理任务，所以大部分推理引擎都会在硬件加速里支持 GPU 加速。

在移动端侧，GPU 加速一般采用对接 OpenGL、OpenCL、Vulkan、Apple Metal，我们以 LiteRT 的 GPU 加速为例说明。

使用 LiteRT 推理的时候可以在创建模型的时候指定开启 GPU 加速：

ERT_ASSIGN_OR_RETURN(auto env, Environment::Create({}));
LITERT_ASSIGN_OR_RETURN(auto compiled_model, CompiledModel::Create(env, model, kLiteRtHwAcceleratorGpu));

指定 GPU 支持的算子：

// in compiled_model.cc
if (hardware_accelerators & kLiteRtHwAcceleratorGpu) {
    const char* accelerator_supported_custom_ops[] = {
        "Convolution2DTransposeBias", 
        "MaxPoolingWithArgmax2D",
        "MaxUnpooling2D", "Resampler"
    };
    for (const auto& op_name : accelerator_supported_custom_ops) {
      resolver.AddCustom(op_name, &sStubRegistration);
    }
}

GPU后端会在初始化的时候自动注册，以Android为例的话，删掉其他平台的宏定义：

static constexpr absl::string_view kGpuAcceleratorLibs[] = {
    "libLiteRtGpuAccelerator" SO_EXT,
    #ifdef __ANDROID__
    #if LITERT_HAS_OPENCL_SUPPORT
          "libLiteRtOpenClAccelerator" SO_EXT,
    #endif  // LITERT_HAS_OPENCL_SUPPORT
    #if LITERT_HAS_WEBGPU_SUPPORT
          "libLiteRtWebGpuAccelerator" SO_EXT,
    #endif  // LITERT_HAS_WEBGPU_SUPPORT
    #if LITERT_HAS_VULKAN_SUPPORT
          "libLiteRtVulkanAccelerator" SO_EXT,
    #endif  // LITERT_HAS_VULKAN_SUPPORT
};

注册顺序为: 通用GPU(OpenGL)加速、OpenCL 加速、Web Gpu 加速、Vulkan

  • OpenGL：主要用于2D和3D图形渲染，但是生态、兼容性、普及程度都比较高
  • OpenCL：主要用于通用并行计算，常用于科学计算、机器学习预处理
  • Vulkan：新一代高性能、低开销的图形与计算 API，OpenGL 的新时代替代品

为什么 GPU 计算能力更强，但是推理引擎不直接全都跑在 GPU 上呢？主要因为以下几个因素：

  1. GPU 擅长大批量、规则、计算密集的任务，CPU 则擅长小批量、低延迟、逻辑复杂的任务。如果是小模型推理跑在 GPU 上，启动 GPU 推理需要
  • 将权重和输入从 CPU 内存拷贝到 GPU 显存
  • 编译/加载 OpenGL 着色器或者 OpenCL 内核
  • 同步 CPU-GPU

如果模型比较小，这些开销可能超过计算本身需要的时间。

  1. 在手机设备上，GPU 共享设备内存，根本没有独立的显存，带宽比较低。频繁在 CPU 和 GPU 之间拷贝数据会抵消计算加速的收益。
  2. GPU 的功耗显著高于 CPU
  3. 兼容性和碎片化问题比较严重，不同设备在 OpenGL、OpenCL、Vulkan的支持程度不同
  4. GPU 开发和调试难度高过 GPU，比较复杂

综合这几个角度，进行 NPU 加速或者专注于 CPU 推理优化会比进行 GPU 加速更好。

NPU加速

NPU（Necural Processing Unit，神经网络处理单元）是一种专为人工智能设计的硬件加速单元。它通过高度定制化的架构，针对神经网络中的常见操作，例如矩阵乘法、卷积、激活函数等进行优化，从而在能效比方面远超CPU 和 GPU。

在相同的手机硬件上，同一个 AI 任务在 CPU、GPU、NPU 上的运行对比如下：

维度CPUGPUNPU
架构特点少核（4-8），高单线程性能，强控制逻辑多核（数百数千），擅长并行计算专用AI核心，无通用指令集
推理速度慢，主要依赖软件库中等快，硬件支持运算
能耗中高，负载高容易发热高非常低
内存带宽频繁访问内存高带宽显存，但是在手机上共用内存片上 SRAM + 数据流架构，减少外部内存访问
精度支持FP32/FP64，效率低支持FP16/INT8支持INT8/INT4/FP16,部分支持混合精度

从能效比方面看，NPU会是移动端模型推理的首选。

市面上常见的 NPU 主要由芯片厂商、手机 SoC(System on a Chip，片上系统) 设计公司和 AI 芯片企业开发，并且会有配套的厂商 SDK，常见包括

  • 高通骁龙芯片的 AI Engine Direct（QNN）
  • 联发科天玑芯片的 NeuroPilot
  • 华为达芬奇架构麒麟系列芯片的 CANN
  • 苹果的Apple Neural Engine(ANE)

下面以高通 QNN、华为 CANN 为例描述 NPU 对接。

QNN

理解 QNN 之前，需要先了解一下高通的 Qualcomm Snapdragon Neural Processing SDK (aka “SNPE”)

SNPE 的工程流程如 SNPE 文档里的解释图，和其他平台的端侧推理框架类似，把一个 Pytorch 训练的模型经过训练、优化，转换为高通 SNPE 的 .dlc 模型文件，通过 SPNE SDK 加载模型和推理。

SNPE 需要将模型转换为专用的 DLC 格式，且仅支持将整个模型部署到单一后端（如 CPU、GPU 或 DSP），无法实现跨硬件的协同计算。为充分发挥骁龙平台异构 AI 硬件（CPU/DSP/NPU）的性能与能效优势，高通推出了 QNN。QNN 提供统一的硬件抽象层，支持算子级别的细粒度调度，可将模型的不同部分动态分配至最适合的计算单元，实现真正的异构计算。

高通文档的架构图详细描述了这一设计思路：

QNN 通过 convert 工具把 Pytorch、TFLite、Tensorflow、ONNX的模型转为 QNN 的格式。QNN 转换器前端把模型转成中间表示，中间表示包含计算图和算子定义，量化后 QNN 转换器后端会把中间表示转为最终的 QnnModel API 调用(cpp文件)。

我们可以使用这个 so 模型进行推理，但是初始化时间会比较长，因为存在一个图构建的编译过程，编译完成后会将模型图结构和上下文以二进制形式缓存到本地存储，也可以用 qnn-context-binary-generator 提前离线编译加快首次运行速度。

CANN

华为 CANN 包括 ATC 模型转换工具和 acl API 库。

ATC模型在模型转换过程种会进行算子调度优化、权重数据重排、内存使用优化等具体操作。

  • 开源框架模型通过 Parser 解析后转为 IR Graph
  • IR Graph经过一些列优化，转成 .om 模型
  • 通过 acl API 加载 om 模型并执行推理

acl API 分为 cpp 版本和 Python 版本，在端侧运行可以使用 cpp 版本。模型加载、推理代码如下：

// 假设以下变量已在类/函数中声明：
// uint32_t modelId_;
// void* modelMemPtr_ = nullptr;
// void* modelWeightPtr_ = nullptr;
// aclmdlDesc modelDesc_ = nullptr;
// aclmdlDataset* input_ = nullptr;
// aclmdlDataset* output_ = nullptr;

const char* omModelPath = "../model/xx.om";

// Step 1: 查询模型所需内存
size_t modelMemSize_ = 0, modelWeightSize_ = 0;
aclError ret = aclmdlQuerySize(omModelPath, &modelMemSize_, &modelWeightSize_);

// Step 2: 分配模型内存
ret = aclrtMalloc(&modelMemPtr_, modelMemSize_, ACL_MEM_MALLOC_HUGE_FIRST);

ret = aclrtMalloc(&modelWeightPtr_, modelWeightSize_, ACL_MEM_MALLOC_HUGE_FIRST);

// Step 3: 加载模型
ret = aclmdlLoadFromFileWithMem(omModelPath, &modelId_, modelMemPtr_, modelMemSize_, modelWeightPtr_, modelWeightSize_);

// Step 4: 获取模型描述
modelDesc_ = aclmdlCreateDesc();
ret = aclmdlGetDesc(modelDesc_, modelId_);

// Step 5: 准备输入
size_t modelInputSize = aclmdlGetInputSizeByIndex(modelDesc_, 0);
void* modelInputBuffer = nullptr;
ret = aclrtMalloc(&modelInputBuffer, modelInputSize, ACL_MEM_MALLOC_HUGE_FIRST);

ret = aclrtMemcpy(modelInputBuffer, modelInputSize, inputBuff, modelInputSize, ACL_MEMCPY_HOST_TO_DEVICE);

input_ = aclmdlCreateDataset();
aclDataBuffer* inputData = aclCreateDataBuffer(modelInputBuffer, modelInputSize);
ret = aclmdlAddDatasetBuffer(input_, inputData);

// Step 6: 准备输出
output_ = aclmdlCreateDataset();

size_t outputSize = aclmdlGetNumOutputs(modelDesc_);
for (size_t i = 0; i < outputSize; ++i) {
    size_t buffer_size = aclmdlGetOutputSizeByIndex(modelDesc_, i);
    void* outputBuffer = nullptr;
    ret = aclrtMalloc(&outputBuffer, buffer_size, ACL_MEM_MALLOC_HUGE_FIRST);
    
    aclDataBuffer* outputData = aclCreateDataBuffer(outputBuffer, buffer_size);
    
    ret = aclmdlAddDatasetBuffer(output_, outputData);
}

// Step 7: 执行推理
ret = aclmdlExecute(modelId_, input_, output_);

size_t outBufferSize = aclmdlGetOutputSizeByIndex(modelDesc_, 0);
void* hostOutput = malloc(outBufferSize);
aclrtMemcpy(hostOutput, outBufferSize,
            aclGetDataBufferAddr(aclmdlGetDatasetBuffer(output_, 0)),
            outBufferSize, ACL_MEMCPY_DEVICE_TO_HOST);
free(inputBuff)

Android NNAPI

Android NNAPI 是 Androd8.1之后Google推出的硬件加速方案。NNAPI 提供了一套C API，实现模型推理图构建和执行。

NNAPI 定义了一批通用算子、模型->编译(构造计算图)->执行计算的流程、 厂商驱动层接口等标准。各个硬件厂商基于 NNAPI 实现自己的驱动程序，包括算子实现、图优化等。对于缺少驱动的设备，会降级到 CPU 上执行计算。

在 Android NDK sample 的代码仓库里面，也有一个 NNAPI 的调用示例：

github.com/android/ndk…

其中，CreateCompiledModel 实现了模型计算图的构造过程，Compute 实现了计算过程。不过实际使用中，应用开发者不需要写这个，这个具体过程一般会由 TF Lite 之类的推理框架来完成。

目前 NNAPI 已经弃用，在 Android15 上已经正式不再使用。

那么 Google 为什么要搞 NNAPI 呢？因为在AI刚开始爆发的时代，硬件加速器非常多，而 Android 设备又非常碎片化，没有办法统一访问这些不同的硬件，Google 推出 NNAPI 是希望屏蔽这种硬件上的碎片化，提供一个标准 API，实现“一次开发，处处加速”。

但是现实往往跟不上理想，在现实中，厂商的实现质量参差不齐，有些只支持部分算子，有些性能实现的还不如 CPU。并且在大量硬件设备面前，应用开发者无法知道不满足预期的表现究竟是模型问题，还是硬件加速算子不支持，但是 NNAPI 没有相应的调试工具，导致黑盒开发，难以进行下去。从厂商的角度，厂商更愿意推广自己的 NPU 加速方案，不愿意跟在 Google 后面做驱动适配。毕竟如果完全听 Google 的，模型推理标准会被 Google 定义，自家硬件独有优势也无法得到体现。

端侧大语言模型

端侧大模型推理优化

大语言模型因为参数权重多，模型文件大，所以大语言模型在端侧部署推理和一般的深度学习模型有一些区别。这些区别主要体现在下面 2 个方面：

  • 大语言模型资源约束较大，LLM 的权重本身就可能超过端侧设备可用内存大小。
  • 大语言模型推理需要专门的优化，包括模型量化、KV Cache管理

大语言模型通常以自回归的方式生成文本，给定一段输入提示词 prompt，模型逐个预测下一个 token，直到完成完整响应。

大模型推理的第一阶段为 prefill（预填充），模型一次性处理整个 prompt，计算所有 token 的上下文表示，并初始化生成过程所需的状态。

大模型推理的第二阶段为 decode，模型以自回归的方式逐个生成输出 token，每一步都根据此前所有 token 的上下文进行单 token 预测。

:::tips 主流大模型基于 Transformer 架构自注意里机制进行推理，自注意力机制工作的时候，每个词都会生成3个关键向量：Q（Query），K（Key），V（Value）

  • Q向量表示：我想要寻找什么信息
  • K向量表示：我能提供什么信息，可以被谁匹配
  • V向量表示：我实际是什么内容

所以在自回归过程中，理解整个句子的语义，都需要依赖前面每个词语的KV向量。

:::

在大语言模型（LLM）的自回归文本生成过程中，模型需要反复计算已生成 token 的上下文表示。如果不做优化，每生成一个新 token，都要重新计算整个历史序列的注意力，导致大量重复计算。

所以大模型推理引入了 KV Cache 机制

  • 在 Prefill 阶段，模型计算输入 prompt 中每个 token 对应的 Key（K）和 Value（V）向量，并将它们缓存起来。
  • 在后续 Decode 阶段，每生成一个新 token，只需计算当前 token 的 Q（Query），然后与缓存中所有历史 K/V 进行注意力计算，无需重复计算已处理 token 的 K/V。

但代价是：KV Cache 会占用大量内存（尤其在长序列或多 batch 场景下），因此也成为端侧推理优化的重点对象。整个大模型推理过程简化为下图：

llama.cpp

llama.cpp 是一个由 Georgi Gerganov 开发的开源项目，旨在用纯 C/C++ 实现对 Llama 系列大语言模型的高效推理，无需依赖 GPU 或深度学习框架。llama.cpp 兼具高性能、轻量级和跨平台特性，支持多种操作系统和硬件架构。在模型支持方面，llama.cpp 不光支持 Llama 系列，还支持 GPT、Qwen2.5等大模型，也支持一些多模态大模型，例如 LLaVA1.5、Qwen2-VL。

llama.cpp的工程架构如下：

核心包括了

  • 推理库：实现模型解码、生成、上下文管理等核心推理逻辑
  • ggml张量计算库：底层计算引擎，提供张量操作和内存管理
  • 模型格式转换工具：将模型统一转换为 GGUF 格式

推理流程

我们看下使用 llama.cpp 在 Android 端调用推理的过程，下面直接放cpp代码，llama.cpp的examples里面有封装好的Kotlin实现。

  1. 初始化、加载模型
// 初始化环境
ggml_backend_load_all_from_path(path_to_backend);
llama_backend_init();

// 加载GGUF模型
const auto *model_path = env->GetStringUTFChars(jmodel_path, 0);
auto *model = llama_model_load_from_file(model_path, model_params);

// 准备阶段，初始化上下文等对象
auto *context = init_context(g_model);
g_context = context;
g_batch = llama_batch_init(BATCH_SIZE, 0, 1);
g_chat_templates = common_chat_templates_init(g_model, "");
g_sampler = new_sampler(DEFAULT_SAMPLER_TEMP);
  1. prefill流程
// 系统提示词的 prompt
const auto *system_prompt = env->GetStringUTFChars(jsystem_prompt, nullptr);
std::string formatted_system_prompt(system_prompt);
// 模版化
if (has_chat_template) {
    formatted_system_prompt = chat_add_and_format(ROLE_SYSTEM, system_prompt);
}
// 分词
const auto system_tokens = common_tokenize(g_context, formatted_system_prompt, has_chat_template, has_chat_template);
decode_tokens_in_batches(g_context, g_batch, system_tokens, current_position);

// 用户输入词的 prompt
const auto *const user_prompt = env->GetStringUTFChars(juser_prompt, nullptr);
std::string formatted_user_prompt(user_prompt);
const bool has_chat_template = common_chat_templates_was_explicit(g_chat_templates.get());
if (has_chat_template) {
    formatted_user_prompt = chat_add_and_format(ROLE_USER, user_prompt);
}
auto user_tokens = common_tokenize(g_context, formatted_user_prompt, has_chat_template, has_chat_template);
decode_tokens_in_batches(g_context, g_batch, user_tokens, current_position, true)

prefill的时候，chat_add_and_format会一路调用到chat.cpp 的 common_chat_templates_apply_jinja 里，这里面会根据模型文件里面存储的元数据来区分不同模型的模板拼接输入的 prompt，例如可以找到Qwen3 coder和xiaomi MiMo的：

  1. decode自回归过程
while (true) {
    generateNextToken()?.let { utf8token ->
        if (utf8token.isNotEmpty()) emit(utf8token)
    } ?: break
}

generateNextToken 实现逻辑：

// 从之前的推理结果获取下一个最可能的token的id
const auto new_token_id = common_sampler_sample(g_sampler, g_context, -1);
common_sampler_accept(g_sampler, new_token_id, true);
// 为下一次decode做输入准备
common_batch_clear(g_batch);
common_batch_add(g_batch, new_token_id, current_position, {0}, true);
// decode
llama_decode(g_context, g_batch);
current_position++;
// 获取token对应字符串并追加到输出结果
auto new_token_chars = common_token_to_piece(g_context, new_token_id);
cached_token_chars += new_token_chars;

llama.cpp 的核心推理过程在于 llama_context::decode，核心流程如下图：

kV Cache优化

为了解决 ggml 张量库一个 DAG 图算子节点数量有上限，庞大的 token 导致计算内存暴涨等问题，llama.cpp 会把 token 输入进一步拆分为微处理批次，每一个微处理批次处理一部分 token。每个微处理批次独立进行图构建、算子计算。同一个批推理的微批处理共享 KV Cache。

llama.cpp 的 KV Cache 根据模型类型采用不同策略实现 C++ 对象的多态调用：

  • llama_kv_cache：标准大语言模型的 KV Cache，连续内存块存储所有缓存
  • llama_kv_cache_iswa：滑动注意力窗口，保存内存恒定大小的 KV Cache，新 token 加入会丢弃老 token
  • llama_memory_recurrent：支持非 Transformer 的循环式语言模型存储推理过程的隐藏状态，本文不深究此类型
  • llama_memory_hybrid：llama_kv_cache 和 llama_memory_recurrent 混合管理

并行计算

llama.cpp 采用异步计算和同步采样的设计，异步计算保证了高吞吐量和计算性能，同步采样则确保了推理的 logits 结果的正确性。

llama.cpp底层自己实现了线程池结构：

struct ggml_threadpool {
    ggml_mutex_t mutex;  // 互斥锁
    ggml_cond_t  cond; // 条件变量

    struct ggml_cgraph * cgraph;
    struct ggml_cplan  * cplan;

    // 同步原语
    atomic_int n_graph; // 图计数
    atomic_int GGML_CACHE_ALIGN n_barrier;
    atomic_int GGML_CACHE_ALIGN n_barrier_passed;
    atomic_int GGML_CACHE_ALIGN current_chunk;

    atomic_bool stop; // 停止
    atomic_bool pause; // 暂停
    atomic_int  abort; // 中断

    struct ggml_compute_state * workers; // 工作线程状态
    int n_threads; // 线程数量
    int32_t prio; // 优先级
    uint32_t poll; // 轮训级别
    
    enum ggml_status ec;
}

在 ggml 图计算的时候，会设置计算的线程池：

ggml_status llama_context::graph_compute(
    ggml_cgraph * gf,
    bool   batched
) {
    int n_threads = batched ? cparams.n_threads_batch : cparams.n_threads;
    ggml_threadpool_t tp = batched ? threadpool_batch : threadpool;
    if (backend_cpu != nullptr) {
        auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_cpu));
        auto * set_threadpool_fn = (decltype(ggml_backend_cpu_set_threadpool) *) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cpu_set_threadpool");
        if (set_threadpool_fn) {
            set_threadpool_fn(backend_cpu, tp);
        }
    }
    for (const auto & set_n_threads_fn : set_n_threads_fns) {
        set_n_threads_fn.second(set_n_threads_fn.first, n_threads);
    }
    // 提交到线程池执行图计算
    auto status = ggml_backend_sched_graph_compute_async(sched.get(), gf);
    return status;
}

LiteRT-LM

LiteRT-LM 是 Google 基于 LiteRT 开发的一个框架，LiteRT-LM 把多个 LiteRT 模型与预处理和后处理组件通过 Piepline 框架组合在一起。他把大语言模型推理过程的分词、prefill、decode 等都显式拆成组件，并且提供状态管理、内存复用等功能。LiteRT-LM 发布于2025年6月，截止2025年12月，仍然处于 alpha 阶段。支持 Gemma3-1B、phi-4-mini、qwen2.5-1.5b等模型。支持CPU、GPU、NPU等多种后端。

架构设计

LiteRT-LM有几个核心API：

  • Engine：负责加载模型及其资源（如分词器）、创建会话。
  • Conversation：表示与 LLM 一次有状态的对话，封装了 Session 并处理复杂的数据任务，包括维护上下文、管理工具定义、多模态数据预处理和 Jinja Prompt 模板，且各实例相互独立，支持并发。
  • Session：提供会话的底层控制，让用户可以手动管理 prefill、decode、多模态数据预处理等细节。

LiteRT-LM 框架设计分为五层

  • 接口层：提供 CLI、Kotlin、C++等API
  • 核心运行时：包括 Engine、Session等
  • 执行后端：LLM 执行器、音视频执行器等
  • 组件层：可复用的机器学习组件，包括分词器、采样器等
  • 基础层：LiteRT框架能力

架构图如下：

在 LiteRT-LM 的 pipeline 里面，定义了如下的组件，pipeline 即是大模型推理里关键推理阶段的定义者，也是会话层和执行器的中间层：

  • Prefill：处理输入提示，填充 KV Cache
  • Decode：自回归生成文本
  • DecodeStreaming：流式生成文本
  • DecodeCustomSampling：自定义采样器解码
  • DecodeCustomSamplingStreaming：自定义采样器解码并且流失生成文本
  • ScoreCustomSampling：文本评分，计算目标文本的对数概率

推理过程

prefill 会提交在线程池运行，并同步等待结果。当前线程只负责任务调度，不负责具体计算：

RETURN_IF_ERROR(worker_thread_pool_.Schedule(
    [this, preprocessed_contents = std::move(preprocessed_contents),
    &status]() 
{
    status = this->PrefillInternal(preprocessed_contents,
    /*wait_for_completion=*/true);
}));
RETURN_IF_ERROR(worker_thread_pool_.WaitUntilDone(Engine::kDefaultTimeout));


// PrefillInternal
absl::Status SessionBasic::PrefillInternal(
    const std::vector<InputData>& preprocessed_contents,
    bool wait_for_completion) {
  ASSIGN_OR_RETURN(ExecutorInputs inputs,ProcessAndCombineContents(preprocessed_contents));
  ASSIGN_OR_RETURN(
      last_prefill_token_id_,
      Prefill(executor_, inputs, wait_for_completion, benchmark_info_)
    );
  return absl::OkStatus();
}

这里会使用 executor_ 执行 prefill 任务，executor_ 类型为 LlmExecutor，执行器有几个子类实现：

  • LlmLiteRtCompiledModelExecutorStatic：静态 shape 模型执行器
  • LlmLiteRtCompiledModelExecutorDynamic：动态 shape 模型执行器
  • LlmLiteRtNpuCompiledModelExecutor：NPU 专用执行器

具体使用哪个执行器，由使用方自己决定传入枚举值。

decode 逻辑类似于 prefill，也是通过上述的执行器执行 decode 过程。

为了应对超长 token 序列的情况，LiteRT-LM 使用工作组将一个长序列拆分为多个子序列处理任务，优化了内存使用效率。

ids = ids.subspan(kTokenIndexToReduce * input_length, input_length);
ASSIGN_OR_RETURN(auto work_groups, GetOptimizedPrefillWorkGroups(
    prefill_signature_map_, ids.size())
);
for (int i = 0; i < work_groups.size(); ++i) {
    const auto& prefill_signature = work_groups[i].first;
    int prefill_length = work_groups[i].second;
    if (!prefill_input_buffers_.contains(prefill_signature)) {
        prefill_input_buffers_[prefill_signature] = {};
        CreatePrefillInputBuffers(
          prefill_signature, 
          prefill_length, prefill_length,
          prefill_input_buffers_[prefill_signature]
        );
    }
    bool async = i < work_groups.size() - 1 || !params.GetWaitForCompletion();
    PrefillInternal(
        prefill_signature, 
        prefill_input_buffers_[prefill_signature],
        ids.subspan(/*pos=*/0, prefill_length),
        async
    );
    ids = ids.subspan(/*pos=*/prefill_length); // 切换到下一个工作组
}
KV Cache与内存管理

与 CPU 不同，GPU 执行的是大规模并行线程，他的内存模型不保证对同一内存地址的并发读写操作具有原子性或者顺序性。如果一个内核正在写入新的 KV Cache，另一个内核同时又读取该区域去计算注意力，那么就会产生数据竞争，出现读写不一致的情况。如果使用 GPU 编程的同步工具进行同步，则会降低 GPU 的并行计算能力，降低吞吐量。

在 LlmLiteRtCompiledModelExecutorStatic 的 Prefill 中，设置了2个KV Cache 缓冲区，来解决 GPU 后端无法同时读写同一段缓冲区的问题。

input_kv_cache_buffers_ = &kv_cache_buffers_1_;	
output_kv_cache_buffers_ = &kv_cache_buffers_2_;

初始化输入缓冲区：

memset(prefill_input_pos_ptr, 0, prefill_input_pos_size);
// 初始化注意力掩码
if (signatures_.input_attn_mask.has_value()) {
    InitializeAttentionMask(
        prefill_input_buffers[signatures_.input_attn_mask.value()],
        use_fp16_precision_
    );
}

模型执行的时候复制缓冲区：

absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer> input_buffers;
// prefill 阶段的prompt
for (const auto& [input_name, input_buffer] : prefill_input_buffers) {
    LITERT_ASSIGN_OR_RETURN(auto input_buffer_dup, input_buffer.Duplicate());
    input_buffers[input_name] = std::move(input_buffer_dup);
}
// 当前轮次的kv cache
for (const auto& [input_name, input_buffer] : *input_kv_cache_buffers_) {
    LITERT_ASSIGN_OR_RETURN(auto input_buffer_dup, input_buffer.Duplicate());
    input_buffers[input_name] = std::move(input_buffer_dup);
}

// 输出缓冲区
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer> output_buffers;
for (const auto& [output_name, output_buffer] : *output_kv_cache_buffers_) {
    LITERT_ASSIGN_OR_RETURN(auto output_buffer_dup, output_buffer.Duplicate());
    output_buffer_dup.ClearEvent();
    output_buffers[output_name] = std::move(output_buffer_dup);
}
// 执行模型
compiled_model_.Run(prefill_signature, input_buffers, output_buffers));
if (!gpu_optimized_single_buffer_cache_) {
    std::swap(input_kv_cache_buffers_, output_kv_cache_buffers_);
}

这里可以看到两个缓冲区会进行交替，每次KV Cache更新后，写入输出缓冲区，然后交换，输出缓冲区作为下一个输入缓冲区。双缓冲区具有下面几个优势：

  • 避免数据竞争，输入输出是独立 buffer，计算时不会读写
  • 交换内存指针，实现了零拷贝，性能高

在 LlmLiteRtCompiledModelExecutorDynamic 的 Prefill 中，首次 prefill，初始化 KV Cache 内存：

if (kv_cache_buffers_1_.empty()) {  
  kv_length = prefill_length;  // 动态解析shape大小
  // 首次预填充，分配 KV cache 缓冲区  
  for (const auto& k_cache_input_name : key_cache_input_names_) {  
    RETURN_IF_ERROR(ResolveDynamicShape(model_, 
        compiled_model_, 
        "prefill",  
        k_cache_input_name, prefill_length)
    );  
    LITERT_ASSIGN_OR_RETURN(  
        auto input_buffer,  
        compiled_model_.CreateInputBuffer("prefill", k_cache_input_name)
      );  
    kv_cache_buffers_1_[k_cache_input_name] = std::move(input_buffer);  
  }  
}

KV Cache 内存扩展，当现有容量不足的时候，计算需要的额外容量，解析动态shape的长度，并扩展缓冲区，因为动态 shape 扩容比较频繁，所以只支持单缓冲区：

int free_kv_entries = kv_length - step_and_token.step;
if (prefill_length > free_kv_entries) {
    int new_kv_seq_len = kv_length + prefill_length;
    int entries_to_add = new_kv_seq_len - kv_length; //计算额外容量
    for (const auto& k_cache_input_name : key_cache_input_names_) {
        ResolveDynamicShape(model_, 
            compiled_model_, 
            "prefill",
            k_cache_input_name,
            new_kv_seq_len
        );
        // 扩展缓冲区
        ASSIGN_OR_RETURN(
            kv_cache_buffers_1_[k_cache_input_name],
            ResizeKVCacheTensorBuffer(env_, kv_cache_buffers_1_[k_cache_input_name],key_dynamic_dim_index_, entries_to_add)
        );
    }
}

高度封装方案

Google 还提供了一些集成程度比较高的解决方案。例如 MediaPipe 和 ML Kit。

MediaPipe

MeidaPipe 提供了一套库和工具，针对一些通用的机器学习场景进行快速接入。这些库和资源提供了

  • MediaPipe Tasks：用于部署解决方案的跨平台 API
  • MediaPipe 模型：预训练的模型

开发者可以使用 MediaPipe 提供的 .tflite 模型，也可以自己的 .tflite 模型，只需要输入向量的 shape 保持兼容即可。MediaPipe 提供了 Model Marker 模型制作工具，这个工具会移除模型的最后几层并进行重新训练，并使用开发者提供的数据重新构建这些层，也提供了 Studio 网页来体验各个任务的效果。MediaPipe 整体组成如下：

MediaPipe 目前支持的几大类型任务为：

  • 生成式 AI：大语言模型推断，图片生成，支持检索增强生成、函数调用
  • 视觉任务：对象检测、图片分类、图片分割等十余种任务
  • 文本任务：文本分类、文本嵌入、语言检测
  • 音频任务：音频分类

从使用的API和模型转换的文档看，MediaPipe 底层使用的技术就是 LiteRT 框架。

ML Kit

ML Kit是 Google 推出的一款移动端 SDK，封装了一系列端侧的 AI 功能。主要包括

  • 生成式AI：摘要、校对、重写和图片描述
  • 视觉任务：OCR、人脸检测、文档扫描等
  • 自然语言：语言识别、翻译、智能回复等

在 Android 中，模型文件可以通过 gms 下载，模型会存储在 Google Play 服务，不计入应用的存储空间，也不会额外增加应用的大小（SDK仍然会增大大小），当模型有新版本发布的时候会通过 gms 自动更新。只有依赖这个模型的所有应用都被卸载，Google Play 服务才会把这个模型从存储里删除。

ML Kit 非常适合海外 Google Play 上架的 Android 应用。

端模型资源部署

Google Play for On-device AI

一个端侧模型的大小有可能非常的大，不可能全都打到 APP 包内，只能通过动态下发。不同模型根据不同条件下载更新，例如不同的 SOC 下发不同的模型。整个模型资源下发也是一整套工程化任务。

在 Android 端设备里，Google 为了减小开发者这部分工程化成本，顺应时代发展推出了 Google Play for On-device AI 库，这个库融合了 App Bundle 和 Google Play 的分发优势，可以自定义模型的分发规则。

这是一个 .tflite 模型通过 Google Play 下发的流程图：

On-device AI 支持在 XML 文件里面进行设备定位。定位属性包括

  • 系统芯片
  • 设备型号
  • 设备RAM
  • 系统功能，主要是一些硬件功能

下载成功后通过 AiPackManager API 获取模型路径：

AiPackLocation aiPackPath = aiPackManager.getPackLocation(aiPack);

if (aiPackPath == null) {
    return null;
}

String aiAssetsFolderPath = aiPackPath.assetsPath();
String aiAssetPath = FilenameUtils.concat(aiAssetsFolderPath, relativeAiAssetPath);

Google 设计一套这样的分发体系，可以帮助 Android 开发者在模型部署上节省很多工作，并且在全球分发的时候具备优势：

  • 隐私合规：通过 Google Play 分发，更容易过审，增加用户信任
  • 免费，无额外 CDN 带宽成本
  • 提升下载性能和可靠性：Google Play 全球拥有优化过的 CDN 网络

目前 On-device AI 处在 Beta 阶段，海外运营 APP 的开发者可以对此方案保持关注。

avatar
文章
阅读
粉丝