Scaling LLM Test-Time Compute with Mobile NPU on Smartphones

  • 文章标题: 在智能手机上利用移动NPU扩展LLM测试时计算
  • 作者/机构: Zixu Hao (清华大学), Jianyu Wei (中国科学技术大学), Tuowei Wang (清华大学), Minxing Huang (清华大学), Huiqiang Jiang (微软研究院), Shiqi Jiang (微软研究院), Ting Cao (清华大学智能产业研究院), Ju Ren (清华大学)

A1 主要贡献

本文旨在解决在移动设备上部署大型语言模型(LLM)时面临的困境:小型模型性能不足,而大型模型资源消耗过大。

核心问题与机遇:
研究发现,在典型的LLM推理过程中,移动神经处理单元(NPU)的计算资源,特别是其矩阵乘法单元,并未得到充分利用。具体来说,在解码阶段,通用的矩阵乘法(GEMM)操作通常会退化为矩阵向量乘法(GEMV),导致为大矩阵块优化的硬件单元利用率低下,造成计算能力的浪费。

研究目标与核心思想:
本文提出利用这种被浪费的计算能力,在移动NPU上应用并行的测试时扩展(test-time scaling)技术,以增强小型LLM的性能。测试时扩展方法通过生成多个候选路径并从中选择最佳样本来提升模型表现,这会增加推理时的计算量,正好可以利用NPU中闲置的计算资源。

面临的挑战:
在移动NPU上高效实现测试时扩展面临两大硬件挑战:
1. 精度问题 (Precision): 移动NPU最初为粗粒度量化模型设计,缺乏对现代LLM在低比特部署中至关重要的细粒度组量化(fine-grained group quantization)的原生硬件支持。使用传统的每通道量化会导致模型在推理任务上性能严重下降。
2. 效率问题 (Efficiency): NPU虽然擅长矩阵乘法,但其通用向量单元的计算吞吐量和内存带宽有限。许多非矩阵运算(如Softmax中的指数计算和反量化过程)必须在向量单元上运行,成为性能瓶颈。

创新点与主要贡献:
为了克服上述挑战,本文设计并实现了一个端到端的LLM推理系统,并提出了两项关键技术:
1. 硬件感知的瓦块量化方案 (Hardware-aware Tile Quantization Scheme): 提出了一种新颖的、同时感知矩阵和向量单元的量化布局。通过在量化前后对权重进行布局转换,将细粒度组量化应用于硬件友好的瓦块(tile),并与NPU的内存访问模式对齐,从而最小化运行时内存访问开销,最大化向量计算利用率。
2. 高效的基于查找表的计算 (Efficient LUT-Based Computation): 将复杂的关键操作,如Softmax中的指数计算和混合精度GEMM中的反量化过程,替换为高效的查找表(LUT)指令,减轻了向量单元的计算瓶颈。

成果总结:
- 系统实现: 设计并实现了一个端到端的LLM推理系统,该系统利用移动NPU支持测试时扩展工作负载,且对专有软件栈的依赖最小。
- 性能提升: 在高通骁龙平台上,所提出的技术为混合精度GEMM带来了高达19.0倍的加速,为Softmax带来了2.2倍的加速。
- 性能-成本权衡: 证明了使用测试时扩展的小型模型可以达到甚至超过大型模型的准确性,实现了新的性能-成本帕累托前沿,为在移动设备上部署LLM开辟了新的可能性。
- 首次探索: 据我们所知,这是首个在移动设备上利用NPU探索LLM测试时扩展方法的可行性并评估其权衡的工作。


图 1. 两种典型的测试时扩展方法:Best-of-N 和 Beam Search。

A3 背景知识与挑战

2. 背景知识

2.1 LLM测试时计算扩展

并行测试时扩展范式。并行测试时扩展是一种新兴且有效的提升模型准确性的范式,它不修改模型参数,而是在测试时投入更多计算。最简单的方法是多数投票和自洽性(self-consistency)【【3,Large language monkeys: Scaling inference compute with repeated sampling+2024+arXiv】, 【61,Self-consistency improves chain of thought reasoning in language models+2023+arXiv】】,它们从多组生成的样本中选出最一致的答案。对于有可验证结果的数学或编程问题,以及有奖励模型的领域(如结果奖励模型,Outcome Reward Models),可以从已完成的样本集中选择得分最高的选项,这种策略被称为Best-of-N【54,Scaling llm test-time compute optimally can be more effective than scaling model parameters+2024+arXiv】。通过前瞻性展开(lookahead rollouts),类似蒙特卡洛树搜索(MCTS)的方法可以从部分生成的序列中选择最优路径,从而衍生出直接对中间结果评分的过程奖励模型(Process Reward Models, PRMs)【【14,Making language models better reasoners with step-aware verifier+2023+arXiv】, 【45,Skywork-o1 open series+2024+https://huggingface.co/ Skywork】, 【60,Math-shepherd: Verify and reinforce llms step-by-step without human annotations+2023+arXiv】, 【74,Processbench: Identifying process errors in mathematical reasoning+2024+arXiv】】。在PRM的辅助下,无需前瞻的步级束搜索(step-level Beam Search)【【54,Scaling llm test-time compute optimally can be more effective than scaling model parameters+2024+arXiv】, 【63,Inference scaling laws: An empirical analysis of compute-optimal inference for problem-solving with language models+2024+arXiv】】可以动态丢弃低质量的生成路径,以平衡探索和利用。图1展示了两种流行的测试时扩展方法的算法。

2.2 神经处理单元 (NPU)

NPU的兴起与架构。随着AI工作负载的增长,现代SoC越来越多地集成NPU来加速神经网络推理【【10,Intel npu acceleration library+2024+https://intel.github. io/intel-npu-acceleration-library/npu.html】, 【41,Qualocmm® hexagon™ npu+2023+IEEE Hot Chips 35 Symposium (HCS)】, 【52,Amd xdna™ npu in ryzen™ ai processors+2024+IEEE Micro】】。NPU的特点是专门加速低精度、计算密集型的核心神经网络操作(如GEMM),在提供极高计算吞吐量的同时保持良好的能效。

“向量+矩阵”组合架构与执行模型。一种被广泛采用的NPU架构是“向量+矩阵”的组合,其中矩阵单元加速矩阵乘法和卷积等操作,而向量单元处理通用计算,如归一化和复杂的激活函数。知名的例子包括高通的Hexagon NPU【41,Qualocmm® hexagon™ npu+2023+IEEE Hot Chips 35 Symposium (HCS)】、华为的Ascend NPU【32,Ascend: a scalable and unified architecture for ubiquitous deep neural network computing: Industry track paper+2021+IEEE International Symposium on High-Performance Computer Architecture (HPCA)】、AMD的XDNA NPU【52,Amd xdna™ npu in ryzen™ ai processors+2024+IEEE Micro】、Intel NPU【10,Intel npu acceleration library+2024+https://intel.github. io/intel-npu-acceleration-library/npu.html】和Intel的Gaudi HPU【11,Intel gaudi 3 ai accelerator white paper+2025+https://www.intel.com/content/www/us/en/content-details/ 817486/intel-gaudi-3-ai-accelerator-white-paper.html】都采用了这类架构。这类NPU在硬件执行模型上与常见的GPU有显著不同。如图2所示,在GPU的SIMT模型中,不同线程可以独立执行分支、内存访问和计算;而在NPU的基于SIMD的执行模型中,单个线程操作大型向量或矩阵数据块。在硬件层面,NPU通常采用较少的硬件线程,并使用VLIW架构来减少控制逻辑开销。与GPU相比,NPU牺牲了编程灵活性和易用性,以换取更高的执行效率和能效。

图 2. (a) GPU的SIMT执行模型与 (b) NPU的SIMD执行模型的比较。

3. 动机与挑战

本节首先介绍移动NPU的一些关键特性,然后分析利用NPU闲置算力实现测试时扩展工作负载的机会与挑战。

3.1 高通Hexagon NPU

Hexagon NPU的代表性。高通骁龙SoC上的Hexagon NPU因其典型的架构、广泛的应用和相对易于获得的SDK,成为移动NPU的代表。因此,我们用它来展示移动NPU的核心特性。

3.1.1 编程接口

编程接口与限制。编程高通Hexagon NPU的主要方式是通过高通AI引擎Direct【51,Qualcomm® ai engine direct sdk+2025+https://developer.qualcomm.com/software/qualcomm-ai-enginedirect-sdk】(常被称为QNN),这是一个专有的、闭源的DNN推理框架。尽 管Hexagon SDK中提供了完整的Hexagon NPU LLVM工具链,但在大多数情况下,开发者无法定制高性能的底层内核,主要是因为矩阵单元的指令集未公开。我们通过对二进制库中未文档化的指令进行逆向工程,得以利用FP16矩阵单元。

3.1.2 架构


图 3. Hexagon NPU 架构。

计算单元。Hexagon NPU采用了典型的“向量+矩阵”混合架构。其向量和矩阵单元分别命名为HVX(Hexagon Vector eXtension)和HMX(Hexagon Matrix eXtension)。Hexagon NPU包含6到8个用于逻辑控制的标量VLIW硬件线程。所有的向量或矩阵指令都由标量核心中的四个VLIW槽之一发出。HVX单元上下文包含32个宽度为1024位的向量寄存器,此类单元的数量在4到6个之间。HMX单元的数量推断为1或2个。

内存子系统。Hexagon NPU包括一个共享的1 MiB L2缓存和8 MiB的TCM(紧耦合内存),后者是一段软件管理的片上内存。HVX可以从L2缓存或TCM读取数据。向量的分散/聚集(scatter/gather)操作和所有HMX指令只能访问TCM。数据可以通过l2fetch指令和DMA机制分别从DDR内存加载到L2缓存和TCM中。两者都支持1D或2D张量数据的异步传输。

HMX单元。Hexagon NPU强大的矩阵乘法能力源于HMX组件。根据高通的说法,HMX单元支持多种精度,包括INT4、INT8、INT16和FP16【41,Qualocmm® hexagon™ npu+2023+IEEE Hot Chips 35 Symposium (HCS)】。以下介绍主要基于FP16 HMX,相关信息来源于逆向工程、Hexagon SDK、QNN SDK以及高通的公开信息。

HMX操作的数据单元。HMX操作的基本数据单元是瓦块(tile),每个瓦块包含一个特定大小的矩阵。对于FP16 HMX,一个瓦块尺寸为32×32,占用2 KiB空间。HMX单元可以从TCM中加载多个权重内存和激活内存的瓦块。在对每对矩阵瓦块执行矩阵乘法后,它会将结果累加到一个内部累加器中。最后,它输出一个与累加器对应的瓦块。同时,HMX单元可以独立地对输出瓦块的每个通道(列)进行缩放和加偏置。


图 4. (a) FP16 HMX瓦块的内存布局。每个瓦块对应一个32 × 32的矩阵,占用2048字节。每两行进行置换,其布局与转置后的2 × 32子矩阵相同。(b) 基于HMX的GEMM的整体内存布局。权重瓦块按列主序排列,因为硬件在瓦块级别执行内积。

FP16 HMX瓦块的内存布局。FP16 HMX瓦块具有特殊的内存布局,如图4(a)所示。输入和输出瓦块都遵循此布局。构建此布局的一种典型方法是使用HVX指令对原始矩阵的每两个相邻行执行跨道(cross-lane)洗牌操作。

3.2 机遇:LLM解码期间的闲置矩阵计算

LLM解码阶段的计算特性。在自回归生成过程中,LLM的输入通常只对应一个token,这导致GEMM操作退化为GEMV。例如,一个形状为[1, hidden_dim]的激活矩阵与一个形状为[hidden_dim, proj_dim]的权重矩阵相乘。在使用FP16 HMX的情况下,每个计算瓦块的有效尺寸为[1,32]×[32,32]。由于硬件计算的基本单元是32×32的瓦块,输入激活瓦块中的31行并没有对应实际有用的内容,导致矩阵单元利用率低下,计算能力被浪费。

测试时扩展的并行性优势。与此同时,一些测试时扩展算法可以通过增加生成期间的计算来获得更好的生成质量,包括自洽性、Best-of-N和束搜索等并行采样方法。它们的特点是使用大于1的批处理大小来探索多个生成路径,并使用某些方式(例如外部验证器)来选择更优的生成路径。图5展示了使用Best-of-N进行测试时扩展的一个例子。随着生成预算(即解码阶段的最大批处理大小)的增加,模型在MATH500数据集上的准确性显著提高。

核心提议。基于以上观察,我们提议在移动NPU上运行LLM的测试时扩展工作负载。通过这种方式,可以有效利用在常规LLM生成过程中被浪费的NPU计算能力。理论上,解码开销不会显著增加,而模型的生成质量可以在运行时得到提升,且无需修改模型权重。

图 5. 使用两个模型进行测试时扩展的示例。随着生成预算的增加,MATH500上的准确率得到提升。

3.3 挑战

理论与实践的差距。尽管利用移动NPU进行测试时扩展在理论上是可行的,但高效的实现面临着许多硬件挑战。我们将这些挑战总结如下。

精度不足。尽管HMX单元支持FP16 GEMM,但在资源受限的设备上部署FP16模型仍然不切实际,使得量化模型成为典型的替代方案。包括HMX在内的大多数移动NPU中的矩阵单元最初是为加速采用粗粒度量化方案(如per-tensor或per-channel量化)的整数化DNN模型而设计的。一个具体的例子是,Hexagon NPU缺乏对现代LLM至关重要的细粒度量化方法的原生硬件支持。这一限制也反映在软件栈中:QNN仅支持per-tensor或per-channel的权重量化。直接对LLM权重应用粗粒度的低比特量化会导致显著的精度下降。

精度下降的实证。如表1所示,Llama 3.2 1B-Instruct模型在QNN的per-channel量化和AWQ per-group 4-bit量化(均为W4A16设置)下的准确性结果表明,per-channel量化模型在具有挑战性的数学推理任务中性能严重下降。不幸的是,由于测试时扩展方法正应用于此类任务,QNN所能达到的基线准确性甚至无法满足性能扩展的最低要求。

表 1. Llama3.2-1B-Instruct在不同实现下的性能比较。QNN的量化严重损害了模型的推理能力。

通用计算和内存带宽薄弱。在缺乏原生硬件支持细粒度组量化的情况下,一种常见的方法是依赖通用计算单元来处理这类计算。然而,我们发现NPU内部的通用向量单元与专用矩阵单元在计算和内存访问能力上存在显著差距。我们使用一个1024×1024×1024的GEMM操作,在Hexagon V75 NPU上测量了HVX和HMX的FP16 GEMM性能,所有输入和输出都驻留在片上TCM中以反映硬件的峰值性能。如表2所示,矩阵单元的FP16 GEMM吞吐量高达12 TFLOPS,比单个向量线程高出300多倍。在内存带宽方面,专用的DMA引擎从DDR实现了超过60 GB/s的读取带宽,而向量单元通过核心数据路径的内存读取带宽仍低于30 GB/s。然而,DMA提供的高带宽仅限于大型、规则的1D或2D数据块,无法高效处理小型或不规则的内存访问。这些观察结果凸显出,向量单元的通用计算和内存带宽不足以跟上专用矩阵单元的计算吞吐量,这为在细粒度量化下实现高性能混合精度GEMM内核带来了重大挑战。

表 2. HVX和HMX单元的性能指标,包括FP16 GEMM计算能力(GFLOPs)和内存读取带宽(GB/s)。

A2 方法细节

4. 设计概述

系统设计目标。我们提出了一个专为移动NPU设计并针对测试时扩展工作负载优化的LLM推理系统。

精度与计算策略。为了应对准确性挑战,我们对主要权重采用4位细粒度组量化,同时保持激活值为浮点数。在运行时,我们动态地将权重即时反量化为浮点值,利用NPU强大的FP16矩阵计算能力来高效支持测试时扩展任务。


表 2. HVX和HMX单元的性能指标,包括FP16 GEMM计算能力(GFLOPs)和内存读取带宽(GB/s)。

应对通用计算瓶颈的核心策略。对于向量处理单元在内存带宽和计算吞吐量上表现有限的不可避免的通用计算,我们的核心策略包括:
- 采用硬件感知的离线设计,以最小化运行时计算开销;
- 充分利用SIMD向量单元的内在能力,以弥合专用硬件和灵活软件需求之间的差距。

具体技术。具体来说,我们引入了以下技术:

硬件感知的细粒度瓦块量化方案

我们提出了一种新颖的量化布局,它在细粒度的矩形瓦块中执行组量化,而不是沿累加轴分组的传统方法。为了与矩阵和向量单元的内存访问模式对齐,我们引入了一个离线流水线,包括权重预量化转换、量化和后量化转换。这增强了运行时内存访问的连续性,并消除了不必要的计算开销。

高效的基于查找表的计算

对于更复杂的运行时操作,我们利用向量单元的查找表(LUT)指令和广义LUT机制来替代复杂的转换逻辑。这种方法加速了测试时扩展工作负载中的关键瓶颈操作,包括混合精度GEMM内的反量化和Attention中的Softmax操作。

5. 系统设计

5.1 硬件感知的细粒度瓦块量化方案

细粒度量化的必要性。已有工作【【30,Quantization meets reasoning: Exploring llm low-bit quantization degradation for mathematical reasoning+2025+arXiv】, 【36,Quantization hurts reasoning? an empirical study on quantized reasoning models+2025+arXiv】】表明,量化误差在数学推理等挑战性任务中会显著降低模型性能。然而,由于设备上严格的资源限制,全精度模型不可行,使得细粒度量化成为保持准确性的关键。

移动NPU上的挑战。不幸的是,在移动NPU上实现高效的基于反量化的GEMM内核在细粒度量化下会带来巨大的系统挑战。我们识别出两个主要问题:
- 矩阵单元期望的权重布局与传统组量化布局之间的不匹配;
- 由小组尺寸引起的宽向量寄存器的次优利用。

解决方案概述。为了克服这些限制,我们提出了一种新颖的瓦块量化方案,包含两个组成部分:
- 一种基于瓦块的量化策略,旨在与矩阵单元固有的数据布局对齐;
- 一种后量化权重置换方法,以最大化向量单元处理能力的利用率。

5.1.1 瓦块-组量化

传统量化GEMM的权重布局。在传统的量化GEMM中,权重矩阵通常以列主序布局存储,这与CPU上基于向量点积的矩阵乘法操作(如llama.cpp CPU后端)相一致。权重沿列维度被划分为连续的量化组——通常大小为32。在每个组内,值被量化,得到的整型权重及其对应的缩放和零点参数以交错方式存储在内存中,保持了矩阵原始的列式顺序。

NPU上的布局不匹配问题。然而,在具有特殊矩阵单元的NPU上,传统的组布局常常与硬件要求不一致。如图6所示,在传统布局中连续的元素在片上TCM中变得分散。对于SIMD向量单元,这种非顺序访问模式是有问题的。尽管现代向量引擎提供gather/scatter操作来缓解分散访问,但这些操作仍然代价高昂。简单地转置权重矩阵并不能解决不匹配问题,因为矩阵单元期望的复杂多级数据布局仍然会导致非连续的内存访问。

图 6. 量化组布局与HMX瓦块布局不匹配的简化示意图。

硬件友好的量化流程。为了解决这个问题,我们首先将权重置换成矩阵单元期望的布局,然后逐组应用四舍五入量化。对于32的组大小,这种方法实际上是在2×16的瓦块单元中执行组量化。鉴于典型模型中的预训练权重近似遵循零均值高斯分布,在这些重塑的瓦块组内进行量化与传统分组相比,并不会显著改变每个组内的统计特性。因此,由此产生的量化误差仍然是可比的。

具体实现。具体来说,我们在量化前根据图4所示的布局排列权重,该布局分层结构为两级:外层是瓦块的列主序排列,以匹配矩阵单元的瓦块级内积操作;内层是每个瓦块内每两行的洗牌。然后,我们按新的内存顺序对权重进行逐组量化。

5.1.2 为宽向量访问合并量化组

默认的AoS布局及其问题。默认情况下,量化权重以结构数组(AoS)布局存储。以Q4_0对称量化为例,每组32个元素由16字节的INT4量化值和2字节的FP16缩放值组成,量化值和缩放值在内存中交错存储。由于NPU架构上的内存访问严重依赖于软件管理的本地1D或2D预取,我们避免使用数组结构(SoA)布局(其中量化值和缩放值位于分开的大型连续数组中),以便更好地与硬件偏好的访问模式对齐。

细粒度量化与向量处理粒度的不匹配。然而,细粒度量化组引入了与原生向量处理粒度的不匹配:单个量化组太小,无法填满一个128字节宽的向量寄存器。访问这样的小组需要多次内存操作或额外的指令来合并来自多个寄存器的数据,导致内存带宽使用效率低下和计算开销。

解决方案:合并为超级组。为了解决这个问题,我们将8个量化组合并成一个更大的超级组,并重组其内容,使得来自256个连续元素的INT4值恰好占据一个完整的HVX寄存器。这个过程如图7所示。

5.2 基于查找表的计算 (LUT-Based Computations)

利用LUT减少计算开销。鉴于向量单元的通用计算性能有限,我们提出使用广义查找表(LUT)指令来替代复杂计算,从而减少指令数量和计算开销。基于LUT的计算对于加速测试时扩展工作负载中的关键操作特别有效,例如Softmax中的指数函数和反量化过程。

5.2.1 通过向量Gather实现快速Softmax

Attention成为瓶颈。测试时扩展方法通常会增加采样并行度,导致更大的批处理大小和更长的上下文长度。我们分析了这些扩展因素在生成过程中对基于Transformer的LLM中主要算子的影响:

图 7. 将8个细粒度量化组重组为一个超级块。INT4量化值正好装入一个向量寄存器。

图 8. Hexagon NPU上FlashAttention的延迟分解。我们使用Qwen2.5-1.5B,提示长度设置为4096。
- GEMM。基于前述NPU硬件特性,在测试时扩展工作负载中适度增加批处理大小不会显著增加GEMM延迟。此外,GEMM延迟与上下文长度无关。
- 其他操作。对于激活函数、LayerNorm、残差加法和RoPE等算子,虽然其计算开销大致与输入大小成正比,但由于其计算量和内存访问量较小,我们忽略了它们的影响。
- Attention。Attention的理论计算复杂度随批处理大小和上下文长度扩展,使其在测试时扩展场景中成为潜在的性能瓶颈。

Softmax在Attention中的主导地位。我们使用FP16 HMX在Hexagon NPU上实现了FlashAttention【12,FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, July 2023.】并测量了在4096提示长度下不同输入批处理大小(查询长度)下的延迟构成,如图8所示。结果表明,矩阵乘法对整体延迟的贡献很小,而随着查询长度的增加,Softmax主导了Attention的执行时间。

指数计算是Softmax的瓶颈。我们的分析表明,片上Softmax的主要瓶颈在于指数计算,该计算必须应用于Θ($N_q$ × $N_{kv}$)个元素。更糟的是,这些昂贵的指数操作必须在HVX上执行,而HVX缺乏对特殊数学函数的专用硬件支持。按照惯例,我们将exp替换为exp2,并在$1/\sqrt{d}$缩放因子中吸收系数$log_2e$。对于一个分解为整数部分$i$和分数部分$f$的输入元素$x$,2^x通过泰勒级数多项式展开来近似,而$i$则直接加到2^i的IEEE-754表示的指数域。然而,多项式求值涉及顺序依赖,限制了VLIW架构下的指令级并行性。

基于LUT的指数计算方案。为了缓解指数计算瓶颈,我们探索用预计算的查找表(LUT)替代显式指数计算。HVX提供了vgather指令,可以将TCM中分散位置的值聚集到一个连续的128字节TCM区域。虽然vgather可以实现大型LUT,但用LUT实现exp仍然具有挑战性:为32位浮点数存储$2^{32}$个元素是不切实际的。此外,vgather本身会引入相当大的延迟——在Hexagon V75上为24到48个指令包,因此必须最小化其使用。

实用LUT方案设计。为了实现实用的基于LUT的exp,我们设计了以下方法。首先,我们在整个FlashAttention中广泛使用FP16,片上计算过程如算法1所示。矩阵Q、K、V和向量m、l存储为16位浮点数,exp计算的输入和输出均为16位浮点数。特别是,FP16 HMX内部使用更高精度的浮点数进行累加,并且我们将元素上转换为32位精度以进行关键操作,如矩阵S的行式求和。

LUT大小与寻址优化。使用16位输入和输出将LUT限制在65536个条目,需要128 KiB的存储空间,这可以放入TCM中。vgather的一个变体支持在一条指令中收集64个2字节元素,最大地址偏移为65536字节。然而,65536个FP16条目占用128 KiB,使得一半的条目无法通过直接寻址访问。为了解决这个问题,我们利用了安全softmax【44,Online normalizer calculation for softmax+2018+arXiv】的特性,该特性通过减去行最大值$m_i$来确保exp的所有输入都是非正数。因此,我们只存储$x \le 0$的值,从而得到一个有32768个条目(64 KiB)的LUT。在基于LUT的exp计算期间,我们忽略FP16输入的最高有效位(符号位),并将输入左移一位以生成vgather所需的字节偏移。

LUT的开销。LUT在系统初始化期间预先计算,在模型推理期间不引入额外开销。它在TCM中占据一个固定的64 KiB区域,仅占总TCM容量的64KiB/8MiB ≈ 0.8%,因此对其他操作的TCM可用性影响极小。

图 9. 通过查表将INT4量化值转换为FP16数值。

5.2.2 以LUT为中心的高效反量化

高效反量化设计。运行时的HVX反量化需要精心设计以避免额外开销。我们提出了一种基于HVX查找表指令的高效反量化过程。vlut16指令能够对源向量寄存器中的每个8位索引在一个包含16个元素的表中进行查找。每个输入字节被转换为一个16位的值,因此vlut16会产生一对寄存器。

通过查表快速转换INT4到FP16。使用vlut16指令,我们直接将4位量化值转换为[-8, 7]范围的FP16值(对于Q4_0量化方案),避免了传统的掩码-解包-转换指令序列。图9展示了两种方法的比较。对于V79之前的Hexagon NPU,所有HVX浮点操作产生的结果都采用一种称为qfloat的内部格式,需要额外的指令才能转换回标准的IEEE-754格式。使用查表消除了这些开销。这种以LUT为中心的设计可以通过简单调整表内容来轻松支持不同的4位编码方案(例如FP4, NF4【13,Qlora: Efficient finetuning of quantized llms+2023+arXiv】, llama.cpp中使用的IQ4_NL)。

通过查表广播缩放因子。一个128字节的HVX寄存器可以容纳两个大小为32的FP16量化组。因此,传统方法是将标量缩放因子广播到整个向量寄存器,然后连接两个寄存器以便与量化值进行后续乘法。然而,通过使用四个组的缩放因子作为LUT内容并应用预定义的常量索引,我们仅用一条vlut16指令就可以实现四个组的缩放因子的广播。

6. 实现

系统实现概述。我们的推理系统基于llama.cpp【16,llama.cpp: Inference of meta’s llama model (and others) in pure c/c++.+2025+https://github.com/ggml-org/llama.cpp】实现,增加了大约7000行C/C++和内联汇编代码。我们使 用Hexagon SDK(版本6.0.0.2)中的LLVM工具链为Hexagon NPU生成代码。我们特别指出,我们的系统不依赖高通的QNN,从而避免了不灵活的静态固定形状计算图。

算法 1:我们FP16 FlashAttention的片上计算(省略了不同的头)

\text{Input: Head dimension } d, \text{ Number of Query tiles } T_q, \text{ Number of KV tiles } T_{kv}, \text{ Query tile size } B_q, \text{ KV tile size } B_{kv}. \\ \text{Input: Matrices } Q_i (\text{FP32}) \in \mathbb{R}^{B_q \times d}, K_j, V_j (\text{FP16}) \in \mathbb{R}^{B_{kv} \times d} \\ 1. \text{ Initialize } O_i^{(0)} = (\mathbb{0}) \in \mathbb{R}^{B_q \times d} (\text{FP16}), m_i^{(0)} = (-\infty) \in \mathbb{R}^{B_q} (\text{FP16}), l_i^{(0)} = (\mathbb{0}) \in \mathbb{R}^{B_q} (\text{FP16}). \\ 2. S_{ij}^{(t)} = \text{MatMul}(Q_i, K_j^\top, \text{AccuType=FP32}) \in \mathbb{R}^{B_q \times B_{kv}} (\text{FP16}); \\ 3. m_i^{(t)} = \max(m_i^{(t-1)}, \text{rowmax}(S_{ij}^{(t)})) \in \mathbb{R}^{B_q} (\text{FP16}); \\ 4. P_{ij}^{(t)} = \text{LUT_Exp}(S_{ij}^{(t)} - m_i^{(t)}) \in \mathbb{R}^{B_q \times B_{kv}} (\text{FP16}); \\ 5. O_i^{(t)} = e^{m_i^{(t-1)}-m_i^{(t)}} O_i^{(t-1)} + \text{rowsum}(P_{ij}^{(t)}, \text{AccuType=FP32}) \in \mathbb{R}^{B_q \times d} (\text{FP16}); \\ 6. l_i^{(t)} = e^{m_i^{(t-1)}-m_i^{(t)}} l_i^{(t-1)} + \text{rowsum}(P_{ij}^{(t)}, \text{AccuType=FP32}) \in \mathbb{R}^{B_q} (\text{FP16}); \\ \text{Output: } O_i = \text{diag}(l_i^{(T_{kv})})^{-1} O_i^{(T_{kv})}.

模块化设计。我们的实现主要由两个模块组成:一个是Hexagon NPU的算子库,它被编译成一个独立的Hexagon DSP共享对象;另一个模块与CPU侧的llama.cpp集成。NPU算子库实现了计算内核、电源管理、硬件资源管理和一个计算线程池。我们为llama.cpp添加了一个Hexagon NPU后端,利用rpcmem共享内存作为底层缓冲区类型。rpcmem是内核dmabuf内存的包装器,支持CPU和NPU之间共享物理内存。相关的分配、释放和映射接口由Android系统供应商库中的libcdsprpc.so提供。通过利用共享内存缓冲区,我们不仅消除了不必要的处理器间数据复制,还尽可能复用了现有的内存管理系统。此外,我们能够将尚未在NPU上实现的算子调度到CPU上运行,实现了与上层应用的无缝集成。

通信与同步机制。在后端初始化阶段,我们调用Hexagon SDK的FastRPC【50,fastrpc: Fastrpc library for linux userspace+2025+https://github.com/quic/fastrpc】工具启动远程NPU会话,并初始化一个共享内存区域用于通信。在NPU侧,一个线程持续轮询这个共享内存区域,以接收来自CPU的计算请求。与默认的RPC实现相比,通过共享内存的通信可以有更低的延迟。我们注意到,在CPU向共享内存写入数据后,NPU不会自动使相应区域的缓存失效,因为骁龙SoC上CPU和NPU之间只有单向一致性。因此,我们在NPU轮询前手动清除缓存。对于包含模型激活值的共享缓冲区,也需要进行类似的缓存维护操作 。

A4 实验环境

  • 硬件平台:

    • 移动设备: 实验在一加Ace3 (搭载骁龙8 Gen 2, NPU架构V73)、一加12 (搭载骁龙8 Gen 3, NPU架构V75) 和一加Ace5 Pro (搭载骁龙8 Elite, NPU架构V79) 三款安卓设备上进行。
    • 服务器: 部分准确性测试在配备NVIDIA RTX3090 GPU的服务器上完成。
  • 模型:

    • 主要模型: 选用Qwen 2.5系列(1.5B, 3B)和Llama 3.2系列(1B, 3B),这些是适合在移动设备上实际部署的尺寸。
    • 扩展模型: 在评估性能-成本权衡时,额外使用了Qwen 2.5 7B模型。
    • 模型变体: 数学推理任务使用Qwen 2.5和Llama 3.2的Instruct版本。
    • 奖励模型: Best-of-N搜索和步级束搜索使用Skywork-1.5B-PRM作为结果奖励和过程奖励评分器。
  • 数据集与评估指标:

    • 测试时扩展任务: 在MATH500和GSM8K两个数学推理任务上评估模型的pass@1准确率,统一使用0-shot CoT提示。
    • 通用准确性: 使用llama-perplexity工具评估WinoGrande准确率、MMLU准确率和Wikitext-2困惑度。
  • 软件与配置:

    • 代码实现: 基于llama.cpp,使用C/C++和内联汇编开发,不依赖高通QNN。
    • 对比基线: 主要与llama.cpp的OpenCL后端进行比较,该后端针对骁龙Adreno GPU优化了Q4_0矩阵乘法内核。同时提供了基于QNN的FP16性能数据作为参考。
    • 量化设置:
      • 大多数矩阵采用Q4_0量化方案 (4.5 BPW)。
      • FFN(前馈网络)的下采样矩阵采用Q8_0量化方案 (8.5 BPW),以减少量化误差并保持模型准确性。

A4 实验结果

7.2 整体性能

  • 测试时扩展的准确性-延迟权衡 (图 10):
    • 实验内容: 评估了Best-of-N和Beam Search两种测试时扩展方法在不同模型、硬件和数据集组合下的准确性与平均解码延迟。
    • 实验结果: 实验数据表明,测试时扩展提供了一个性能-成本的权衡空间,并在特定配置下实现了更优的帕累托前沿。例如,在Best-of-N方法中,Qwen2.5 1.5B和3B模型的扩展结果分别超越了3B和7B模型的基线准确性。
    • 结论: 通过利用NPU的计算能力和测试时扩展算法,小型设备端模型在生成质量和推理成本方面有潜力超越大型模型。


图 10. 不同测试时扩展方法在各种数据集和硬件组合上的准确性-延迟权衡。

  • 设备端解码性能 (图 11):
    • 实验内容: 在三款设备上测试了不同模型在不同批处理大小下的端到端解码吞吐量。
    • 实验结果: 随着批处理大小的增加,系统的解码吞吐量显著提升。这是因为HMX单元的闲置算力被有效利用。吞吐量并未实现完美的线性扩展,原因是lm_head(词汇表投影)的计算被放在CPU上执行,当批处理大小增加时,这部分成为瓶颈。
    • 结论: NPU在处理更大批次大小的解码任务时表现出良好的扩展性,但系统性能受限于CPU-NPU混合部署中的瓶颈部分。


图 11. 不同模型在各种批处理大小和硬件设置下的端到端解码吞吐量。

  • 功耗和能耗 (图 12):
    • 实验内容: 在OnePlus 12上测量了LLM解码过程中的功耗和能耗。
    • 实验结果: 1.5B模型功耗随批次增加而上升,但设备总功耗仍在5W以内。3B模型功耗稳定在4.3W左右。能耗的扩展趋势与解码延迟相似。特别地,1.5B模型在批次为8时的解码能耗低于3B模型在批次为1时的能耗,而此时前者的数学任务准确性与后者相当。
    • 结论: 测试时扩展在能效上具有优势,可以用更低的能耗实现与更大模型相当的性能。


图 12. LLM解码阶段的功耗和能耗。

  • 与其他系统的比较 (图 13):
    • 实验内容: 将本系统与基于GPU的实现(llama.cpp-OpenCL)和基于QNN的FP16性能进行比较,涵盖解码和预填充阶段。
    • 实验结果: 在解码阶段,虽然GPU在批次为1时更快,但本NPU系统在更大批次下展现出更高的解码吞吞吐量和更好的扩展性。在预填充阶段,本系统也持续优于GPU实现,并能在某些工作负载下达到与专有QNN相当的性能。
    • 结论: NPU在测试时扩展这类并行计算工作负载中比GPU更具优势,且在预填充阶段也具有竞争力。


图 13. 推理吞吐量比较。

7.3 准确性评估

  • 量化方案 (表 4):
    • 实验内容: 比较了为HMX布局定制的瓦块量化组与传统量化组对Qwen2.5-1.5B模型准确性的影响。
    • 实验结果: 使用本文提出的量化布局的模型在MMLU上准确性略高,在Winogrande和Wikitext PPL上略有下降,但差异远小于量化本身带来的性能损失。
    • 结论: 本文提出的瓦块量化组方案不会导致量化模型准确性显著下降。


表 4. 使用为HMX布局量身定制的瓦块量化组的模型与使用传统量化组的模型的准确性比较。

  • Attention实现 (表 5):
    • 实验内容: 比较了使用本文基于LUT的FP16 Attention与传统FP32 Attention的模型准确性。
    • 实验结果: 将Attention中的非关键部分(累加除外)替换为较低的FP16精度,对模型的端到端准确性没有明显影响。
    • 结论: 本文的FP16 Attention实现能够在不牺牲模型准确性的前提下提升性能。


表 5. 使用我们带有基于LUT的Softmax的F16 FlashAttention的模型与使用传统F32 Attention的模型的准确性比较。

7.4 消融研究

  • Attention中的Softmax (图 14):
    • 实验内容: 在不同Attention工作负载下,比较了使用不同方法计算指数函数exp的片上softmax延迟。
    • 实验结果: 基于LUT的指数计算相比传统的32位浮点exp实现了1.26到2.19倍的加速,相比16位浮点exp最高加速1.60倍。
    • 结论: LUT方法显著加速了Softmax中的指数计算,是Attention优化的关键。


图 14. 我们提出的带有基于LUT的指数计算的F16 Attention的片上softmax消融研究。性能在一加12上测量。

  • 基于反量化的GEMM (图 15):
    • 实验内容: 对GEMM反量化布局的优化进行消融实验,比较了基线(传统布局+scatter操作)、仅应用HMX布局、以及应用全部优化(HMX布局+量化组合并)的效果。
    • 实验结果: 与基线相比,本文方法在不同矩阵尺寸下实现了9.65到19.04倍的加速。应用HMX布局和量化组合并带来了1.82倍到3.45倍的加速。本文方法的性能仅比理论上限(不进行实际反量化)慢27%。
    • 结论: 硬件感知的权重布局和量化组合并是实现高效反量化的核心,其性能已接近理论极限。


图 15. 针对GEMM反量化的优化建议的消融研究。我们在OnePlus 12上测量了GEMV的性能。

7.5 开销与敏感性分析

  • CPU与内存使用 (图 16):
    • 实验内容: 在OnePlus 12上评估了1.5B和3B Qwen2.5模型在解码阶段的CPU利用率和内存消耗。
    • 实验结果: NPU内存使用恒定,总内存消耗对于1.5B和3B模型分别约为1.3 GiB和2.4 GiB。CPU利用率随批次增加而上升,但使用的核心数被限制在4个以内。
    • 结论: 系统资源开销在可控范围内。


图 16. 解码阶段的CPU和内存使用情况。

  • 提示长度的影响 (图 17):
    • 实验内容: 分析了提示长度(从512到4096 tokens)对解码吞吐量的影响。
    • 实验结果: 随着提示长度增加,解码吞吐量呈现温和的下降趋势,但下降幅度相对较小。
    • 结论: 在4096 tokens范围内,提示长度对解码吞吐量的影响有限。


图 17. 提示长度对解码吞吐量的影响。

A7 补充细节

8. 讨论

对其他硬件的普适性。我们认为NPU的“向量+矩阵”架构具有一定的普遍性,并观察到CPU和NPU之间的界限正逐渐模糊。除了NPU,现代CPU也开始集成专用的矩阵乘法单元,如Intel AMX和ARM SME,使其具备了类似的“向量+矩阵”架构。此外,我们注意到现代AI加速器(例如NVIDIA GPU)普遍在通用计算性能和专用低精度矩阵乘法能力之间存在显著差异。尽管具体硬件架构可能不同,我们技术背后的核心思想仍具有广泛的适用性。

系统性能与局限性
- (a) 解码性能:当前系统的解码速度相对受限,主要源于反量化的开销。然而,这并未削弱测试时扩展的有效性。基于QNN的量化GEMM通常仅使用DMA和HMX组件,不引入HVX计算开销。类似T-MAC【62,T-mac: Cpu renaissance via table lookup for low-bit llm deployment on edge+2025+European Conference on Computer Systems】的方法可能在NPU上实现高效的细粒度组量化GEMV,从而加速LLM解码过程。
- (b) 预填充性能:当前系统的预填充性能仍有提升空间。将更多算子卸载到NPU,通过算子融合减少内存访问和通信开销,以及优化矩阵乘法的分块和流水线策略,都有助于提升预填充性能。我们将这些优化留作未来工作。
- (c) 模型大小限制:我们当前的实现受限于旧设备上单个NPU会话的32位地址空间。采用多个NPU会话可能有助于缓解此问题。

并行测试时扩展的应用范围。尽管并行测试时扩展方法目前主要在数学推理任务中占据主导地位,但近期研究【【8,Parallel scaling law for language models+2025+arXiv】, 【17,Reasoning as planning: LLMs as building blocks of a rational agent+2024+The Twelfth International Conference on Learning Representations (ICLR)】, 【22,Metagpt: Meta programming for multi-agent collaborative framework+2023+arXiv】, 【49,Boosting llm reasoning: A crossover of tree of thoughts and retrieve-augmented generation+2023+arXiv】, 【69,Empowering agentic video analytics systems with video language models+2025+arXiv】, 【71,Tree of thoughts: Deliberate problem solving with large language models+2023+Advances in neural information processing systems】】的证据表明,它们可扩展到更广泛的推理和规划领域,显示出巨大的可推广潜力。

9. 相关工作

使用NPU的设备端LLM推理。llm.npu【66,Fast on-device llm inference with npus+2025+Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1】开创性地在NPU上使用per-tensor量化的INT8 GEMM来加速LLM的预填充阶段,并利用CPU辅助处理离群值以保持准确性。HeteroLLM【7,Heterollm: Accelerating large language model inference on mobile socs platform with heterogeneous ai accelerators+2025+arXiv】通过张量分区实现了GPU和NPU之间的协同推理。PowerServe【2,Powerserve+2025+https://github.com/powerserveproject/PowerServe/tree/main】是一个开源推理框架,利用中间ONNX格式实现自定义的量化和浮点计算分区。ShadowAttn【73 ,Dynamic sparse attention on mobile socs+2025+arXiv】利用NPU和CPU/GPU之间的协作来加速稀疏注意力。ExecuTorch【【48,Torchao: Pytorch-native trainingto-serving model optimization+2025+arXiv】, 【58,Executorch: A pytorch platform for on-device deployment+2025+https://github.com/pytorch/executorch】】是一个知名的开源边缘端DNN推理框架,支持SpinQuant【38 ,Spinquant: Llm quantization with learned rotations+2024+arXiv】和QNN后端。以上所有工作都基于高通的Hexagon NPU,并使用闭源的QNN【51,Qualcomm® ai engine direct sdk+2025+https://developer.qualcomm.com/software/qualcomm-ai-enginedirect-sdk】作为后端。来自vivo的研究【【6 ,Edgeinfinite-instruct: Bridging sft-based optimization and npu-level efficiency for edge devices+2025+arXiv】, 【40,Bluelm-v-3b: Algorithm and system co-design for multimodal large language models on mobile devices+2024+arXiv】】利用了联发科的NPU,但由于NeuroPilot SDK不公开,此类研究仍然稀少。

LLM量化。最知名的训练后量化算法包括GPTQ【15,Gptq: Accurate post-training quantization for generative pre-trained transformers+2022+arXiv】和AWQ【34,Awq: Activation-aware weight quantization for on-device llm compression and acceleration+2024+Proceedings of Machine Learning and Systems】,它们执行仅权重量化,且仅需少量校准数据,因此被广泛使用。后续方法如SmoothQuant【65,Smoothquant: Accurate and efficient post-training quantization for large language models+2023+International Conference on Machine Learning】、DuQuant【33,Duquant: Distributing outliers via dual transformation makes stronger quantized llms+2024+Advances in Neural Information Processing Systems】将范围扩大到权重-激活量化,通过开发技术来减轻激活值中棘手的离群值分布问题。最近,出现了如QuaRot【1,Quarot: Outlier-free 4-bit inference in rotated llms+2024+Advances in Neural Information Processing Systems】和SpinQuant【38,Spinquant: Llm quantization with learned rotations+2024+arXiv】等综合性方法,旨在量化所有主要组件——权重、激活值和关键的KV缓存——通常低至4位。这些方法利用旋转变换来在整个模型中创造更适合量化的特征分布。

推测解码。推测解码(Speculative Decoding)【【29,Fast inference from transformers via speculative decoding+2023+International Conference on Machine Learning】, 【39,Improving multi-candidate speculative decoding+2024+arXiv】, 【43,Specinfer: Accelerating large language model serving with tree-based speculative inference and verification+2024+Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】】是一类LLM推理的加速方法,其核心是在一次模型前向传播中验证多个推测的token,以缓解LLM解码的内存瓶颈问题。推测解码有多种扩展变体,其中一些【【19,Hybrid slm and llm for edge-cloud collaborative inference+2024+Proceedings of the Workshop on Edge and Mobile Foundation Models】, 【31,Reward-guided speculative decoding for efficient llm reasoning+2025+arXiv】】不再严格遵循目标模型的分布。理论上,广义的推测解码和测试时扩展方法都属于广义的“生成-然后-验证”框架,我们的系统理论上可以无缝支持这些应用。

A5 结论

本工作展示了利用移动NPU(特别是高通Hexagon NPU)中未被充分利用的计算能力来进行LLM测试时扩展的可行性和有效性。通过设计一个集成了硬件感知的瓦块量化、权重布局优化以及基于LUT的关键算子加速的端到端推理系统,我们证明了经过测试时扩展增强的小型模型可以在准确性和延迟方面超越传统部署的大型模型。这种方法为在资源受限的移动设备上部署高性能语言模型提供了一条新途径,推动了设备端AI效率与能力的帕累托前沿。

方法细节中的引用汇总

  • 第5.1节 (硬件感知的细粒度瓦块量化方案) 中,引用了 [30, 36] 来支撑“量化误差会显著降低模型在数学推理等挑战性任务中的性能”这一观点,从而强调了细粒度量化的必要性。
  • 第5.2.1节 (通过向量Gather实现快速Softmax) 中,引用了 [12] 来表明本文在Hexagon NPU上实现了FlashAttention算法。同时,引用了 [44] 来支持使用“安全softmax”的特性,该特性通过确保softmax输入为非正数,从而将指数计算的LUT大小减半。
  • 第5.2.2节 (以LUT为中心的高效反量化) 中,引用了 [13] (QLoRA) 作为例子,说明本文提出的以LUT为中心的设计可以轻松支持包括NF4在内的不同4位编码方案。
  • 第6节 (实现) 中,引用了 [16] 指出本推理系统是基于开源项目 llama.cpp 实现的。同时,引用了 [50] 来说明系统初始化时调用了Hexagon SDK的 FastRPC 工具来启动远程NPU会话。