LLM新巅峰||北大提出mllm-NPU,赋能1000字/秒的设备端LLM预填充

本文提出了mllm-NPU，这是第一个有效利用设备上神经处理单元（NPU）卸载的LLM推理系统。mllm-NPU的主要设计目标是减少预填充延迟和能耗。它针对主流的仅解码器的变压器架构的LLM（例如LlaMA、GPT等）。关键思想是最大限度地在移动NPU上执行预填充，以加速整数计算，同时保持在CPU/GPU上执行必要的浮点操作以保持准确性。为了克服上述挑战并提高NPU卸载效率，mllm-NPU在三个层面上重构了提示和模型：（1）在提示层面：mllm-NPU将可变长度的提示分割成多个固定大小的块，同时保持数据依赖性；（2）在张量层面：mllm-NPU识别并提取重要的异常值以在CPU/GPU上运行；（3）在块层面：mllm-NPU根据硬件亲和性和对准确性的敏感度将变压器块调度到CPU/GPU和NPU。

​​https://arxiv.org/pdf/2407.05858​​

mllm-NPU系统设计概述

mllm-NPU是一个创新的大型语言模型(LLM)推理系统，专为移动设备上的神经处理单元(NPU)优化。该系统的主要设计目标是减少预填充(prefill)阶段的延迟和能耗，同时保持推理精度。mllm-NPU通过在算法和系统设计上的协同优化，有效地解决了LLM架构与现代NPU设计之间的语义差异。

mllm-NPU的工作流程包括准备阶段和执行阶段。在准备阶段，mllm-NPU使用增强的张量级量化算法将LLM量化为W8A8格式，并生成固定长度的分块共享图，以高效处理变长提示。在执行阶段，mllm-NPU将接收到的提示分割成固定大小的块，并按照因果关系处理这些块。这些块图会根据它们的数据格式被分割成子图，并根据硬件亲和性和精度敏感性在CPU/GPU和NPU之间进行调度。

三大创新技术详解

1. 分块共享图

mllm-NPU通过引入分块共享图来应对动态提示长度的挑战。这种方法基于对LLM操作符的两种分类：静态操作符和动态操作符。静态操作符如线性层和层归一化，可以跨不同块共享；而动态操作符如注意力机制，依赖于块长度和序列，不能共享。通过这种方式，mllm-NPU将LLM分解为多个子图，共享子图一次构建和优化，非共享子图则针对不同块单独构建。这种方法显著减少了内存开销，并提高了可扩展性。

2. 影子离群值执行

为了在不牺牲LLM精度的前提下实现NPU友好的张量级激活量化，mllm-NPU采用了影子离群值执行技术。该技术在运行时将具有离群值的激活通道提取到一个更紧凑的张量中，在CPU上执行，并将结果合并回NPU上的原始操作符结果中。这种设计不仅解决了离群值的问题，还优化了内存使用，并减少了CPU和NPU之间的同步开销。

3. 无序子图执行

mllm-NPU采用无序子图执行策略来优化执行效率。通过在块和子图级别进行分区后，LLM子图可以不按块序列顺序执行。这种策略显著扩大了调度空间，最小化了由于CPU/GPU浮点操作引起的执行泡沫。mllm-NPU利用在线启发式算法，在保证正确性的前提下，动态选择执行顺序，从而最大化NPU的利用率并最小化CPU/GPU的工作负载影响。

实现与评估

1. 实现

mllm-NPU的实现涉及多个关键技术，包括分块共享图、阴影异常执行和乱序子图执行。这些技术共同优化了移动NPU的加载效率，显著提高了预填充速度和能效。

• 分块共享图：通过将变长提示分解为多个固定大小的块，并为每个块构建预优化的计算图，mllm-NPU显著减少了图构建和优化的时间。此外，通过共享不依赖于块序列的静态运算符，减少了内存占用，提高了扩展性。
• 阴影异常执行：为了在不牺牲精度的情况下实现对NPU友好的张量激活量化，mllm-NPU采用了阴影异常执行技术。通过在CPU上并行处理稀疏的异常激活通道，mllm-NPU能够在保持NPU计算效率的同时，减少由于异常值引起的量化误差。
• 乱序子图执行：mllm-NPU通过乱序执行子图来优化计算流程，减少由于CPU/GPU上的浮点运算导致的执行泡沫。这种策略扩大了调度空间，使得可以更灵活地管理各个子图的执行顺序，从而最大化NPU的利用率和减少总体执行时间。

2. 评估

mllm-NPU的性能通过与多个基线系统的比较得到验证。实验结果显示，无论是在预填充速度、能耗还是端到端推理延迟方面，mllm-NPU都显著优于现有技术。

• 预填充性能：在1024个令牌的提示长度下，mllm-NPU在Xiaomi 14设备上的预填充速度比使用CPU的llama.cpp快18.17-38.4倍，比使用GPU的TFLite快1.27-2.34倍。
• 能耗：在Redmi K60 Pro设备上，mllm-NPU在1024令牌提示长度下的能耗比使用CPU的llama.cpp低35.63-59.52倍，比使用GPU的TFLite低1.85-4.32倍。
• 端到端性能：在实际的移动应用场景中，mllm-NPU在端到端推理延迟上也展现出了优越性，特别是在UI自动化和自动电子邮件回复的任务中，mllm-NPU的表现远超其他基线系统。

性能分析

mllm-NPU的性能优势主要源于其创新的设计和优化策略。通过在移动NPU上最大化整数运算的执行，同时在CPU/GPU上处理必要的浮点运算，mllm-NPU不仅提高了执行速度，还降低了能耗。

• 内存消耗：尽管mllm-NPU引入了额外的内存开销（主要是由于阴影异常执行技术），但这一开销相对较小，仅占总内存的一小部分。这得益于其高效的内存管理策略，如分块共享图技术，有效减少了不必要的内存占用。
• 精度保持：mllm-NPU在保持高效NPU执行的同时，几乎不损失推理精度。通过动态处理异常值位置，mllm-NPU能够在元素级别上提供比传统量化方法更高的精度，从而保证了模型的输出质量。
• 系统兼容性：mllm-NPU支持标准的LLM格式，并可以与现有的解码引擎兼容，使其能够灵活地集成到各种移动设备和应用中。此外，mllm-NPU的设计允许与其他算法级优化技术结合使用，进一步提升性能。

总之，mllm-NPU通过一系列创新技术，显著提高了移动设备上大型语言模型的推理性能，为实现高效、低能耗的移动设备端LLM推理提供了有效解决方案。

准确性与内存消耗

1. 准确性保持

mllm-NPU系统在实现高效的NPU卸载的同时，保持了较高的准确性。与FP16相比，准确性损失不到1%，这在实际应用中是可接受的。这得益于其创新的技术，如影子离群值执行（shadow outlier execution），它允许在CPU上以浮点精度处理离群值，从而保持了高准确性。此外，mllm-NPU与其他基线模型相比，在保持准确性的同时，显著提高了执行效率。

2. 内存消耗

尽管mllm-NPU引入了一些额外的内存开销，例如为离群值执行加载的浮点权重，但这些开销相对较小，仅占总内存的0.6%至1%。此外，通过共享图（chunk-sharing graph）技术，mllm-NPU显著减少了内存需求，因为它允许在不同的执行块之间共享静态操作符。这种方法不仅减少了内存占用，还提高了内存使用的效率。

本文转载自 ​​AI论文解读​​，作者：柏企