一句话:提出 ProSE 蛋白质设计发现加速器,用异构脉动阵列高效执行 Protein BERT 推理,显著提升速度与能效。
蛋白质语言模型在结构预测、功能注释和药物发现中表现突出,但训练与推理计算成本高,尤其长序列下限制其广泛应用。本文提出 ProSE(Protein Systolic Engine)的架构、微架构与硬件实现,包含定制异构脉动阵列和特殊函数,面向迁移学习模型推理;它将 SIMD 风格计算与脉动阵列结合,跨模型层优化粗粒度操作序列,在提升效率的同时保持通用性。评估表明,ProSE 执行 Protein BERT 推理相比 NVIDIA A100 GPU 最高取得 6.9× 加速和 48× 能效;相比 TPUv3/TPUv2 最高取得 5.5×/12.7× 加速和 173×/249× 能效。
Protein language models have enabled breakthrough approaches to protein structure prediction, function annotation, and drug discovery. A primary limitation to the widespread adoption of these powerful models is the high computational cost associated with the training and inference of these models, especially at longer sequence lengths. We present the architecture, microarchitecture, and hardware implementation of a protein design and discovery accelerator, ProSE (Protein Systolic Engine). ProSE has a collection of custom heterogeneous systolic arrays and special functions that process transfer learning model inferences efficiently. The architecture marries SIMD-style computations with systolic array architectures, optimizing coarse-grained operation sequences across model layers to achieve efficiency without sacrificing generality. ProSE performs Protein BERT inference at up to 6.9× speedup and 48× power efficiency (performance/Watt) compared to one NVIDIA A100 GPU. ProSE achieves up to 5.5 × (12.7×) speedup and 173× (249×) power efficiency compared to TPUv3 (TPUv2).