Prose: The architecture and design of a protein discovery engine

Conference Paper

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).

Full Text

Duke Authors

Cited Authors

  • Robson, E; Xu, C; Wills, LW

Published Date

  • February 28, 2022

Published In

  • International Conference on Architectural Support for Programming Languages and Operating Systems Asplos

Start / End Page

  • 655 - 668

International Standard Book Number 13 (ISBN-13)

  • 9781450392051

Digital Object Identifier (DOI)

  • 10.1145/3503222.3507722

Citation Source

  • Scopus