太初元碁 x 山东大学 | 国产AI芯片赋能分子对接,异构众核架构开启药物研发新时代

一款新药从研发到上市,平均需要耗费10年+时间、投入10亿+美元。如何加快药物研制,缩短研发周期,提升设计药物在实验阶段的成功率等,是药物研发必须面对和解决的核心问题。
太初元碁与山东大学高性能计算团队基于国产异构众核架构达成深度产学研合作,在太初元碁自研的AI芯片上完成UCSF DOCK3.7的全栈适配、移植与性能优化。单机八卡上,每小时可以实现350万+的分子对接。
新药研发之困:
超大规模虚拟筛选的“大海捞针”与算力瓶颈救
药物发现的本质,是一场在浩瀚化学空间中的“大海捞针”。
一个药物靶点(如与疾病相关的蛋白质),需要在数以亿计的类药小分子库中,找到那把能够精准“开锁”的分子钥匙。这个过程被称为分子对接,即在计算机中模拟小分子与蛋白质靶点的结合模式与结合强度,从而筛选出最有可能成药的候选分子。
然而,随着化合物数据库规模的爆炸式增长,分子对接面临严峻的计算挑战。传统的x86 CPU架构在单核性能提升趋缓的背景下,已经难以支撑超大规模虚拟筛选的实时性需求。算力,成为制约新药研发速度的新瓶颈。

图1 算力成为制约新药研发速度的新瓶颈
校企协同攻坚:
太初元碁异构算力,落地经典分子对接引擎
UCSF DOCK3.7因其计算速度快、筛选精度合理,非常适用于药物筛选初期模拟筛选包含多达数十亿小分子的超大规模类药小分子,因此被全球学术界和工业界广泛使用。
太初元碁基于独特的异构众核架构,深度适配UCSF DOCK3.7,完整复现分子对接核心流程,对接精度与原版 UCSF DOCK3.7 保持严格一致。实现从单核到大规模并行的全链路性能释放,能够在合理时间内高效处理数十亿乃至上百亿小分子,加速药物筛选流程。

图2 各阶段研究的化合物量级
全栈深度优化:
四大核心技术,激活国产AI芯片算力潜能
太初元碁与山东大学高性能计算团队,在太初元碁AI芯片上实现从算法逻辑、访存模式、指令压缩、并行框架等四大维度的系统性优化:

图3 全栈深度优化
优化一:算法逻辑优化:从根本上减少无效计算
原始的FindMatches算法在搜索配体-受体结合模式时,会因多次调整距离阈值而产生大量重复的匹配结果,需要额外的去重模块清理。团队重构了算法逻辑,从搜索源头即避免了重复匹配的产生,从而移除了整个去重模块,在不损失任何计算精度的前提下,显著降低了计算量与访存总量。
优化二:DMA 批量访存 + 压缩存储:榨干内存带宽
针对太初元碁AI芯片的片上存储容量特点,团队采用了两步策略:一是用DMA批量传输替代高延迟的全局访存指令,将数据从主存成块搬运至片上存储容量;二是将原本用32位整数表示的原子/sphere索引压缩为8位(因单个配体原子数通常不超过200),使关键索引数组得以完全驻留在片上存储容量中。这些优化使得计算过程中对主存的访问量大幅降低,访存效率显著提升。
优化三:SIMD 向量化:让向量处理单元满负荷运转
充分利用太初元碁AI芯片的SIMD向量处理单元,对计算热点进行手动的向量化优化。通过单指令同时处理16个浮点数,将计算密集型任务的指令数减少至原来的1/16。
优化四:双层 MPI 分布式框架:为多核扩展铺路
面向未来更大规模的虚拟筛选需求,团队设计并实现了基于双层调度的分布式并行框架,有效降低了大规模扩展时的MPI通信瓶颈,并利用聚合I/O带宽实现了高效的并行数据读取。
价值实现:
不止于快,更是国产算力赋能生命科学的里程碑
经过上述的系统性优化,太初元碁与山东大学高性能计算团队取得了一系列成果:
成果一:性能飞跃
在单卡上,优化后的分子对接程序较CPU实现了突破性的性能提升。 在单机八卡上(龙芯3C6000 CPU+太初AI加速卡),每小时可以实现350万+的分子对接。

图4 分子对接示例
成果二:存储与I/O革命
设计的二进制存储格式 + xzip压缩方案,大幅提升数据库存储空间利用率,同时显著缩短文件读取与初始化时间。
成果三:超大规模扩展能力
优化后的程序可高效扩展到太初元碁多个计算核心上,在郑州空港算力集群上实测(龙芯3C6000+太初AI加速卡)扩展效率保持在90%以上。
未来展望:
从“能用”到“好用”,构建HPC+AI的融合生态
此次与山东大学的合作,是太初元碁在HPC+AI领域的一次成功实践。它证明了太初元碁的异构众核架构不仅能够高效支撑AI大模型的训练与推理,同样能在科学计算、生命科学等传统高性能计算领域发挥关键作用。
当每一颗国产芯片都能高效解码生命的奥秘,新药研发的核心桎梏终将被打破。
太初元碁愿与产学研各界伙伴携手,共建开放、协同、繁荣的HPC+AI融合生态,让国产AI算力成为驱动科学发现与产业创新的核心引擎。



