最近做了一个端侧AI的项目，把一个大语言模型部署到手机上运行。最开始，模型在手机上跑一次推理要好几秒，用户体验很差。经过一系列优化，最终把推理时间降到了几百毫秒，实现了从慢到快的转变。

这篇文章就来分享一下这次性能优化的实战经验，从模型压缩、量化、算子优化到硬件加速，聊聊端侧AI推理优化的思路和方法。如果你也在做端侧AI部署，希望这篇文章能帮到你。

端侧AI的挑战

先说说端侧AI面临的挑战。

端侧AI，就是把AI模型部署在手机、平板、IoT设备等终端设备上运行，而不是在云端服务器上运行。端侧AI的好处是：低延迟（不需要网络传输）、隐私保护（数据不上传）、离线可用（没有网络也能用）、节省云端成本。

但端侧AI也面临很多挑战：

第一，算力有限。手机的CPU、GPU、NPU的算力，和服务器上的GPU比起来，差了几个数量级。服务器上的A100有几百TOPS的算力，手机的NPU可能只有几十TOPS。

第二，内存有限。手机的内存一般是8GB、12GB，而且系统和其他应用还要占用一部分，留给AI模型的内存很有限。一个7B的大模型，FP16精度下就要14GB内存，根本装不下。

第三，功耗和发热。手机的散热能力有限，长时间高负载运行会发热降频，影响性能和体验。

第四，电池续航。AI推理是耗电大户，如果推理太频繁或者太耗时，会很快耗尽电池。

所以，端侧AI的核心挑战就是：在有限的算力、内存、功耗下，让模型尽可能快地运行。

优化前的状态

说说我们优化前的状态。

我们的模型是一个7B参数的大语言模型，最初用FP16精度部署，模型大小14GB。在旗舰手机上，第一次推理（加载模型+生成第一个token）要8秒，之后每个token生成要500毫秒。生成一段100字的回答，需要将近一分钟。

这个速度，用户根本没法用。而且，14GB的模型，大部分手机都装不下，只有16GB内存的旗舰手机才能勉强运行。

我们的目标是：模型大小降到4GB以内，第一次推理降到2秒以内，每个token生成降到100毫秒以内。

为了达到这个目标，我们做了一系列优化。下面一个个来说。

优化一：模型量化

第一个优化是模型量化。

量化，就是把模型的权重和激活值从高精度（FP32、FP16）转换成低精度（INT8、INT4），从而减小模型大小、提高推理速度。

我们最开始用的是FP16，模型大小14GB。先试了INT8量化，模型大小降到了7GB，推理速度提升了约40%。但7GB还是太大，大部分手机装不下。

然后我们试了INT4量化，模型大小降到了3.5GB，推理速度又提升了约50%。3.5GB的大小，8GB内存的手机也能运行了。

但INT4量化有一个问题：精度损失。量化之后，模型的生成质量会下降，特别是在复杂推理和长文本生成上。我们做了评测，INT4量化后，模型在一些基准测试上的得分下降了5-10%。

为了减少精度损失，我们用了一些技巧：

第一，混合精度量化。不是所有层都用INT4，敏感的层（比如注意力层的某些部分）用INT8，不敏感的层用INT4。这样在保持模型大小的同时，减少了精度损失。

第二，量化感知训练（QAT）。在训练过程中就引入量化的误差，让模型适应低精度的表示。我们用了量化感知微调，在一个小数据集上微调了几个epoch，精度损失降到了2-3%。

第三，使用更先进的量化算法。比如GPTQ、AWQ这些针对大语言模型的量化算法，能在INT4精度下保持更好的精度。我们最终用了AWQ量化，精度损失控制在2%以内。

量化之后，模型大小从14GB降到了3.5GB，推理速度提升了约2倍。这是最有效的一步优化。

优化二：模型压缩

第二个优化是模型压缩。

量化是减小每个参数的位宽，模型压缩是减少参数的数量。我们用了几种模型压缩的方法。

第一个是剪枝。剪枝就是去掉模型中不重要的参数或者神经元。我们用了结构化剪枝，去掉了一些不重要的注意力头和前馈网络的神经元。剪枝20%之后，模型大小又减小了20%，推理速度提升了约15%，精度损失很小（不到1%）。

第二个是知识蒸馏。知识蒸馏就是用一个大模型（教师模型）来训练一个小模型（学生模型），让小模型学习大模型的输出分布。我们用原始的7B模型作为教师，训练了一个3B的学生模型。3B模型的大小只有1.5GB（INT4），推理速度比7B快了一倍多，精度损失约5%。

第三个是投机解码（Speculative Decoding）。投机解码是用一个小模型（草稿模型）快速生成候选token，然后用大模型（目标模型）一次性验证多个token。如果小模型生成的候选token被大模型认可，就可以一次接受多个token，从而提高生成速度。我们用了一个0.5B的小模型作为草稿模型，7B模型作为目标模型，生成速度提升了约30%。

通过这些压缩方法，我们在保持精度的前提下，进一步减小了模型大小、提升了推理速度。

优化三：推理框架优化

第三个优化是推理框架的优化。

选对推理框架，对端侧AI的性能影响很大。我们最开始用的是通用的推理框架，性能一般。后来换了专门针对端侧优化的推理框架，性能提升了不少。

目前主流的端侧推理框架有：

  • llama.cpp：针对CPU优化的大模型推理框架，支持多种量化格式，跨平台
  • MLC-LLM：基于TVM的大模型推理框架，支持GPU和NPU加速
  • TensorRT-LLM：NVIDIA的推理框架，主要针对NVIDIA GPU
  • Core ML：苹果的机器学习框架，针对iOS设备优化，支持ANE（苹果神经引擎）
  • NNAPI / QNN：安卓的神经网络API，支持NPU加速
  • MediaPipe：谷歌的端侧AI框架，支持多平台

我们针对不同的平台，选择了不同的框架：

  • iOS平台用Core ML，充分利用ANE加速
  • 安卓平台用llama.cpp + Vulkan，利用GPU加速，同时支持NPU通过QNN调用

换了框架之后，推理速度又提升了约30%。特别是在iOS平台上，Core ML + ANE的组合，性能比纯CPU推理快了好几倍。

除了选对框架，我们还对框架做了一些定制化优化：

  • 自定义算子：框架不支持的算子，自己实现优化版本
  • 内存复用：推理过程中，中间张量的内存尽量复用，减少内存占用和分配开销
  • 批处理优化：对于可以批处理的算子，合并多个请求一起处理
  • 预分配内存：推理前预分配所有需要的内存，避免推理过程中动态分配

优化四：KV缓存优化

第四个优化是KV缓存优化。

大语言模型推理的时候，会把之前生成的token的key和value缓存起来（KV缓存），避免重复计算。随着生成的文本越来越长，KV缓存占用的内存也越来越大。

对于长文本生成，KV缓存可能会占用几个GB的内存，成为内存和性能的瓶颈。

我们做了几个优化：

第一，量化KV缓存。把KV缓存也量化成INT8甚至INT4，减小内存占用。量化KV缓存对精度的影响很小，但能显著减少内存占用。我们把KV缓存量化成INT8之后，内存占用减少了一半。

第二，分页注意力（PagedAttention）。借鉴了vLLM的思路，把KV缓存分成固定大小的块（page），像操作系统的虚拟内存一样管理。这样可以减少内存碎片，提高内存利用率，也能更高效地处理多个并发请求。

第三，滑动窗口注意力。对于超长文本，不是缓存所有历史token的KV，而是只缓存最近的N个token。这样可以限制KV缓存的最大大小，避免内存溢出。当然，这会影响模型对长距离依赖的处理能力，需要根据业务场景权衡。

通过这些优化，KV缓存的内存占用减少了60%，长文本生成的速度也提升了不少。

优化五：硬件加速

第五个优化是硬件加速。

现在的手机，除了CPU，还有GPU和NPU（神经网络处理单元）。不同的硬件适合不同的计算任务，合理利用各种硬件，能显著提升推理性能。

CPU的特点是通用性强、延迟低，但算力有限。适合做控制逻辑、预处理、后处理，以及一些不规则的计算。

GPU的特点是并行计算能力强，适合做矩阵乘法、卷积等大规模并行计算。大模型的大部分计算（比如注意力、前馈网络）都适合在GPU上跑。

NPU的特点是专门针对神经网络计算优化，能效比高，适合做大规模的张量计算。现在的旗舰手机NPU算力很强，而且比GPU更省电。

我们的策略是：

  • 模型的主体计算（矩阵乘法、注意力）放在NPU上跑，充分利用NPU的高算力和低功耗
  • NPU不支持的算子，放在GPU上跑
  • 控制逻辑和预处理后处理，放在CPU上跑
  • 不同硬件之间的数据传输，尽量减少，因为传输也有开销

硬件加速之后，推理速度又提升了约40%，而且功耗降低了不少，手机不会那么烫了。

优化六：系统级优化

第六个优化是系统级的优化。

除了模型和推理框架，系统层面也有一些优化空间。

第一个是内存管理。端侧AI对内存很敏感，我们做了精细的内存管理：模型加载的时候用内存映射（mmap），减少内存占用；推理过程中及时释放不需要的内存；和系统协调，避免被系统杀掉。

第二个是线程调度。推理是多线程的，合理设置线程数和线程优先级，能提升性能。线程太少，算力用不满；线程太多，调度开销大。我们根据设备的CPU核心数，动态调整线程数。

第三个是预热和缓存。第一次推理的时候，模型加载和初始化比较慢。我们在应用启动的时候，就提前加载模型、做一次预热推理，这样用户真正使用的时候，就不会觉得慢。同时，把一些常用的计算结果缓存起来，避免重复计算。

第四个是功耗和发热控制。我们监控设备的温度，当温度过高的时候，适当降低推理性能（比如减少线程数、降低频率），避免过热降频。虽然性能降低了一点，但保证了稳定性和用户体验。

优化后的效果

经过这一系列优化，最终的效果是：

  • 模型大小：从14GB降到了3.5GB（INT4量化 + 混合精度）
  • 第一次推理：从8秒降到了1.5秒（模型加载优化 + 预热）
  • 每个token生成：从500毫秒降到了80毫秒（量化 + 框架优化 + 硬件加速）
  • 内存占用：从14GB以上降到了5GB以内（模型压缩 + KV缓存优化）
  • 功耗：降低了约40%（NPU加速 + 功耗控制）

生成一段100字的回答，从原来的将近一分钟，降到了8秒左右。虽然还不如云端API快，但对于端侧来说，这个速度已经可以接受了。

而且，模型可以在8GB内存的手机上运行，覆盖了大部分主流机型。

一些经验和教训

分享一下这次优化的经验和教训。

第一，量化是端侧AI的必经之路。不量化，大模型根本在端侧跑不起来。INT4量化是目前的主流选择，配合好的量化算法（AWQ、GPTQ），精度损失可以接受。

第二，选对推理框架很重要。不同的框架，性能差距很大。不要自己造轮子，用成熟的、活跃的开源框架，在它们的基础上做定制化优化。

第三，硬件加速是关键。纯CPU推理的性能有限，一定要利用GPU和NPU。特别是NPU，能效比很高，是端侧AI的未来。

第四，内存是端侧AI的瓶颈。不仅要关注模型大小，还要关注KV缓存、中间张量的内存占用。精细的内存管理，能让模型在更低配置的设备上运行。

第五，精度和性能要权衡。优化的过程中，精度会有损失，要通过评测来确保损失在可接受的范围内。不要为了追求速度，把模型优化得不能用了。

第六，要在真实设备上测试。模拟器和真机的性能差距很大，不同型号的手机性能也不一样。一定要在多种设备上测试，确保兼容性和性能。

第七，用户体验是最终目标。优化的目的是让用户用得爽，不是为了跑分。要关注用户感知到的延迟（比如第一个token的时间），而不只是理论上的吞吐量。

写在最后

端侧AI是一个很有前景的方向，随着手机算力的提升和模型优化技术的进步，越来越多的AI应用会在端侧运行。

但端侧AI的优化，是一个系统工程，涉及模型、框架、硬件、系统等多个层面。没有银弹，需要一步步来，从最有效的优化开始，逐步提升。

这次从慢到快的优化经历，让我对端侧AI有了更深的理解。也让我意识到，AI的未来不仅在云端，也在端侧。端侧和云端的结合，会是AI应用的主流模式。

如果你也在做端侧AI，欢迎交流。让我们一起，把端侧AI做得更快、更好、更省电。