基于珠峰1.0的llama.cpp矩阵计算优化实践¶

1 前言¶

“珠峰1.0”是中兴微电子发布的首颗领域定制融合处理器，基于Neoverse-N2架构，支持ARMv9指令集，能够有效支撑AI框架的推理任务。该平台可以支持llama.cpp开源大模型推理框架的推理任务。

为了进一步释放该平台性能，砺睿微基于该处理器对开源大模型推理框架llama.cpp中矩阵乘法算子进行了深度优化，本文将对这一实践中用到的优化技术和优化成果进行介绍。

2 优化技术¶

llama.cpp中很多矩阵乘法算子的基础函数命名中都包含“vec_dot”，即表示依赖传统的向量点积实现矩阵乘。而大模型中，参与点积的向量维度一般很大（DeepSeek-V3的典型值是7168），当2个长向量的点积操作进行到很后面时，通常缓存效率并不是很好。而当切换到权重矩阵的下一列时，激活矩阵需要重新回到向量起始位置重新获取数据，此时会产生一个非常大的内存访问的地址跳跃，进一步加剧访问抖动。业界研究发现，可以进一步利用批处理的思想进行**量化布局改造**，以改善这些问题。ARM官方论文 "Highly Optimized Kernels and Fine-Grained Codebooks for LLM Inference on Arm CPUs" 就是基于Q4_0的布局进行重组优化的工作，并在小模型上取得了不错收益。本文工作则将该优化思想迁移到了Q4_K上，以加速更复杂量化类型的更大模型的推理性能。

2.1 传统布局¶

Q4_0的传统布局方式如下图：

基于传统布局的激活向量工作逻辑如下： 1. 逻辑上：激活矩阵的每个Q8_0单元都参与固定对位区的Q4_0单元的对位点积，如上图的对位区1和2。点积顺序必须按照权重矩阵的内存顺序来，依次是黄色列、红色列、绿色列、蓝色列。 2. 实现上：用neon寄存器管理同一列的计算任务，包括缩放因子计算与点积和计算。

传统布局的计算过程存在一定局限性： 1. 相同对位区的计算任务中，由于不同列的Q4_0单元物理不连续，Q8_0的数据无法复用。 2. 计算任务只能一次面向一列，任务并行度不高，neon寄存器的通道利用率也不高。

2.2 重组布局优化¶

下面是Q4_0重组布局的效果：

重组布局的核心变化是在权重矩阵列上打破原有的物理存储顺序，将原来属于相邻多个列的数据从物理不连续改为物理连续。上图中的一个直观体现就是，在传统布局中，同样位置的Q4_0单元在不同色系列中是物理不相邻的，现在通过对每个量化单元内的32个元素进行分段，再将相同位置的其他列的分段部分交错拼接起来，就形成了上图用箭头表示物理存储顺序的新形态。 此时激活向量有新的工作逻辑： 1. 逻辑上：每个对位区中一个Q8_0单元与4个交错排布的Q4_0单元进行点积。 2. 实现上：需要基于neon寄存器宽度，对新的点积工作进行拆分，重点是4个通道分别用于4列数据的计算，包括缩放因子计算和点积和计算。

该布局通过如下两方面改善了传统布局的局限性问题： 1. 同一个对位区的Q8_0单元可以同时参与原先4个不同列的点积计算，计算复用度增长4倍。 2. 任务并行度增加，neon寄存器通道利用率增加。

2.3 基于重组布局进一步优化¶

基于权重矩阵的重组布局这一核心，还有一些其他衍生的优化手段，比如： 1. 激活矩阵也可以进行重组，进一步改善激活矩阵的访存效率。但这点在实现上和权重矩阵的重组有区别，因为后者实质上是全过程静态不变的数据，重组过程可以离线进行，例如在加载模型时进行，不参与推理运行时开销，因此对性能要求不高。而激活矩阵是推理过程中不断迭代的，它的重组操作就对性能非常敏感，建议**利用simd进行高效实现**。 2. 激活矩阵与权重矩阵的计算可以进一步使用arm架构的**矩阵加速指令**来优化。整个llama.cpp的矩阵计算框架分为gemv和gemm，前者形式上是向量和矩阵相乘，但内部实现本质上是向量与向量的乘法，核心是SDOT/UDOT向量点积指令（Armv8.4推出），而后者形式上是矩阵与矩阵相乘，内部实现也是矩阵与矩阵的乘法，核心可以利用更高阶的SMMLA/UMMLA矩阵乘法指令（Armv8.6推出）。

3 模型实测数据¶

本文实验基于珠峰1.0服务器，主要利用llama.cpp自带的llama-batched-bench工具进行性能测试，其结果可以直接反馈提示词处理PP和词元生成TG两个阶段的性能表现。本文源码已经提交社区PR15719，其对比基线可参考commit4795c91。 测试使用的模型是DeepSeek-V3（下载页面）的Q4_K_M模型，该模型是Deepseek 开发并发布的一个强大的混合专家 (MoE) 语言模型，其总参数量达 671B，每个 token 激活参数量为 37B。为了实现高效的推理和经济高效的训练，DeepSeek-V3 采用了多头潜在注意力 (MLA) 和 DeepSeekMoE 架构。模型文件大小 377GB。

上图展示了PP阶段重组布局优化前后的性能对比。纵坐标为吞吐性能（token/s），横坐标为提示词序列数量。随着提示词序列数目增长，PP吞吐会趋于收敛。数据表明，优化后PP吞吐性能平均提升29.2%。由于PP是计算密集阶段，因此矩阵计算性能的提升对其帮助很明显。

上图展示了TG阶段重组布局优化前后的性能对比。TG阶段吞吐性能平均提升4.2%。TG是访存敏感阶段，因此收益相对小。

4 总结¶

本文工作主要基于珠峰1.0服务器，对矩阵计算领域的布局重组技术和向量、矩阵加速指令进行综合运用，对llama.cpp中Q4_K量化模型的运算过程进行深度优化，并在DeepSeek-V3模型上验证了显著的性能提升效果，释放了珠峰平台在AI领域的潜力。