人工智能正在进入科学研究的核心计算环节。从蛋白质候选序列探索到复杂科学模型推理，科研平台既需要模型能力的进步，也需要高效、稳定的计算支持。

面向高水平科技自立自强和科研基础设施建设，让科学模型在国产计算平台上高效运行，是把算力资源转化为科研支撑能力的一项具体工作。围绕这一目标，我们以约14亿参数的ESM3开放模型为对象，在鲲鹏920专业版上开展了矩阵内核、算子融合、数据布局、运行时和多NUMA调度的协同优化。

在FP32、36线程、批量2、序列长度512、8步生成的单NUMA测试中，完成一批生成任务的时间由15.228秒缩短至5.549秒，相对同节点PyTorch原始计算路径加速2.74倍，耗时减少63.56%。八NUMA持续任务队列达到1335.09个最终序列位置/秒，约合每秒完成2.61条长度512的序列。

在FP32设置下，八NUMA持续队列吞吐达到单卡GPU批量8短时测试记录的2.25倍，具体配置与结果见下文。

从科学模型出发，围绕真实计算路径优化

ESM3的推理涉及线性投影、注意力计算、前馈网络以及多种输出任务。模型运行效率不仅取决于矩阵运算速度，也受到中间张量读写、线程同步、调用开销和内存局部性的影响。

我们从完整生成任务入手定位热点，再针对内核、数据流和执行过程展开优化。性能测试覆盖三个层次：独立算子、完整生成和多NUMA持续任务队列，分别观察计算效率、任务延迟和并发产出。

这个过程也揭示了一个实际问题：单个算子变快之后，完整生成未必同步提速。数据复用、缓冲区生命周期和线程等待，都需要与内核计算一起分析。

矩阵内核：针对真实形状组织计算

线性投影承担了模型中的大量计算。针对查询、键和值投影、注意力输出投影以及前馈网络的不同矩阵形状，我们利用矩阵加速单元实现专用内核，并调整分块、加载顺序和线程任务组织。

优化重点包括：

• 根据实际矩阵形状选择微内核与调度块；
• 提高输入和权重的复用，减少重复加载；
• 调整累加范围，减少部分和恢复与写回；
• 控制输出布局转换和最终复制的开销。

此前四类未融合线性投影的独立测试，相对PyTorch获得了1.86至2.15倍加速。在此基础上，进一步通过融合与执行策略优化，将算子收益转化为完整生成任务的提速。

算子融合：让中间结果直接进入下游计算

前馈网络是融合优化的重要对象。分离执行时，上投影输出需要先写入内存，经过激活后，再转换成下投影所需的数据布局。这些步骤会产生中间张量，并带来额外的数据写出和读取。

融合实现采用配对通道计算：在矩阵累加寄存器中完成局部输出块的整个K维累加，随后执行SwiGLU激活，并将结果直接写入下投影所需的打包布局。

如图1所示，融合路径把配对累加、激活和下游布局写入连接起来，省去了完整上投影输出U，以及普通布局激活张量H的物化与重读。

图1 ESM3前馈网络的分离执行与融合执行数据流

激活在对应输出块完成整个K维累加之后执行。图中虚线框标示融合范围；归一化位于范围之外，完整下游A-pack缓冲区和上下游阶段屏障仍然保留，下投影在上游全部完成后启动。

这样的设计先消除不必要的中间数据搬运，同时保留清晰的计算阶段，便于验证正确性和分析性能。

数据布局与调度：平衡输入复用和权重供给

融合改变了数据流，也改变了缓存复用关系。省去中间张量之后，如果输入或权重被反复读取，节省的访存仍可能被抵消。

为此，配对微内核按16行、32对通道组织，结合输入预取与权重供给进行对照测试。下投影采用分段累加、面板间隙，以及规则的矩形任务划分与蛇形遍历，改善访问关系和缓存复用。

调参以完整生成结果为依据，同时观察计算时间、线程等待和内存行为。这样可以识别那些在单算子测试中更快、进入真实模型后却没有收益的配置。

运行时与内存：减少计算之外的开销

随着算子加速，调用包装、重复对象查询、临时分配和同步等待的影响逐渐显现。

项目通过C++调用计划、注意力调用边界精简、残差融合和工作区复用，减少重复准备与调用开销。输入形状、数据布局、线程约束等检查，以及支持范围外的回退机制继续保留。

内存优化同时关注缓冲区生命周期和页面局部性：减少反复回收，建立内容一致的本地运行时与常量副本，并检查页面实际落点，降低跨NUMA访问对性能的影响。

这些措施共同作用于完整任务。最终收益通过端到端对照测量，而不是把各次局部实验的加速比相乘。

以PyTorch为基线，观察逐阶段优化效果

为了观察不同优化环节的作用，实验在同一NUMA上设置了四组对照：

1. PyTorch原始计算路径；
2. 仅启用按需输出头；
3. 启用按需输出头与几何分支捷径；
4. 完整优化版。

纯序列生成只计算任务需要的输出头。几何分支捷径则在输入没有有效结构框架、且满足检查条件时启用。

如图2所示，批量2、长度512的完整生成耗时，从15.228秒逐步降至5.549秒。

图2 相对PyTorch原始计算路径的逐阶段优化效果

以下测试统一采用FP32、36线程和8步序列生成，每组进行2次预热与5次正式计时，各组使用相同的运行时局部性与内存分配策略。

执行配置	批量2、长度512	批量1、长度1024
PyTorch原始计算路径	15.228秒	19.415秒
按需输出头	14.794秒	18.740秒
输出头与几何分支捷径	12.727秒	13.926秒
完整优化版	5.549秒	9.864秒
完整版相对原始路径加速比	2.74倍	1.97倍
完整版相对两项简化后加速比	2.29倍	1.41倍

按需输出头和几何分支捷径减少了任务不需要的计算。在这两项简化之后，完整优化版在批量2、长度512上仍获得2.29倍加速，在批量1、长度1024上获得1.41倍加速，体现了底层内核、算子融合及执行策略的综合收益。

表中的阶段差值反映按这一顺序叠加优化的效果，各项收益存在相互影响。

这轮四组对照共完成40次正式生成，另进行了预热及两组随机种子的正确性检查。51项跨组检查全部通过，覆盖相应对照中的输出头、逐步序列logits、最终输出、序列和随机数状态。正式计时未出现重大缺页，内存保护未触发。

八NUMA持续队列，提升批量任务产出

NUMA是具有本地处理器与内存访问关系的资源域。在多NUMA执行中，每个资源域部署独立模型副本，使用本地资源处理任务。

八NUMA配置下，每个副本使用36线程，共288个工作线程，每副本批量为2。副本完成一批任务后立即领取下一批，减少快副本等待慢副本的空闲时间。

八NUMA持续任务队列达到1335.09个最终序列位置/秒。 两次独立进程组、五组配对队列中，优化版完成1280个正式任务，平均每批生成时间为6.036秒，P99为6.228秒；正式阶段无重大缺页或交换，八副本合计RSS峰值约152.34 GiB。

为进一步观察并发扩展效果，另一次实验对1、2、4、6 NUMA进行了相同节点上的PyTorch对照。

如图3所示，六NUMA吞吐由347.04提高至1009.80个最终序列位置/秒，达到2.91倍。

图3 1至6 NUMA持续任务队列吞吐对照

图3采用每节点36线程、批量2的配置，吞吐按完成的最终序列位置总数除以队列墙钟时间统计，包含任务领取、校验和排空。

其中，2.91倍对应六NUMA同节点对照；1335.09个最终序列位置/秒对应八NUMA持续队列实验，分别反映已测规模下的优化效果与任务产出能力。

FP32设置下，与单卡GPU实测记录对照

以下列出CPU与单卡GPU的FP32实测结果。任务均为长度512、8步序列生成，GPU关闭自动混合精度和TF32。

FP32执行配置	每批完整生成时间	吞吐：最终序列位置/秒
单NUMA优化版，批量2	5.549秒	184.54
单卡GPU，批量2	1.805秒	567.24
单卡GPU，批量8	6.891秒	594.41
八NUMA优化版，每副本批量2	持续任务队列	1335.09

如图4所示，八NUMA持续队列吞吐达到单卡GPU批量8短时FP32测试记录的2.25倍，达到GPU批量2记录的2.35倍。

图4 八NUMA持续队列与单卡GPU短时FP32吞吐记录对照

这一结果体现了CPU多副本并行处理任务的批量产出能力。对于同批量单任务，GPU具有更低延迟：批量2时，单卡GPU的生成速度为单NUMA优化版的3.07倍。单任务响应时间与多副本合计吞吐，分别对应不同的应用需求。

GPU采用单卡批量2和批量8，预热2次后正式计时5次，取批次耗时均值。CPU运行包含按需输出头等优化的序列生成路径，GPU运行官方完整计算路径。

CPU队列计时包含任务领取、校验与排空；GPU采用短时批次生成计时，所选监测窗口未采到热降频标志。吞吐统一以最终序列位置/秒表示，按最终序列位置总数除以耗时计算，不额外乘以生成步数。

上述配置下，八NUMA持续队列与单卡GPU批量8的实测吞吐比为2.25倍。

让国产科研算力形成可复用的应用能力

这次优化从ESM3的实际计算路径出发，逐步解决了矩阵计算、中间数据搬运、缓存复用、线程等待和运行时开销等问题，最终在单NUMA完整生成和多NUMA持续任务产出上获得了可测量的提升。

单NUMA完整生成加速2.74倍，意味着相同任务可以更快完成；八NUMA持续队列每秒完成约2.61条长度512的序列，则为批量候选探索提供了更高的任务处理能力。

这些结果背后，是计算内核、融合数据流、内存布局和系统执行的共同调整。矩阵加速单元的能力需要合适的内核来发挥，内核收益需要有效的数据复用来维持，多核资源则需要合理的调度与内存放置来支撑。

面向AI for Science，国产科研算力建设需要把硬件能力与科学应用的计算特征紧密衔接。从模型热点出发，通过可复核的实验选择内核、布局和执行策略，才能让平台资源转化为研究人员实际可用的推理能力。

持续完善科学模型的软件适配与性能优化，让已有计算资源完成更多候选探索和模型推理任务，是国产科研基础设施走向高效应用的重要一步。

---

测试说明

• 模型： ESM3开放模型，约14亿参数，模型源码与权重版本锁定。
• 框架： PyTorch 2.11.0。
• 精度： CPU采用FP32；GPU采用FP32，关闭自动混合精度和TF32。
• 单NUMA： 36线程，8步生成；测试批量2、长度512，以及批量1、长度1024；每组2次预热、5次正式计时。
• 八NUMA： 每节点36线程、每副本批量2、长度512、8步生成；两次独立进程组、五组配对持续队列。
• 1至6 NUMA： 每节点36线程、批量2，在相同节点集合上对照PyTorch与优化版。
• GPU： 单卡，批量2与批量8，长度512、8步生成；2次预热、5次正式计时。
• 吞吐定义： 完成的最终序列位置总数除以耗时；长度512的完整序列按512个位置计算。
• 结果范围： 本文汇总FP32实验结果，后续BF16实验另行报告。