Attention in SRAM on Tenstorrent Grayskull

当Transformer的自注意力层使用SRAM而非DRAM进行实现时，它们可以实现显著的加速。Tenstorrent Grayskull架构提供了一个大型的SRAM，分布在一个核心网格中。本文提出了一个用于Grayskull的融合内核，通过结合矩阵乘法、注意力分数缩放和Softmax操作，专门利用其大型SRAM。此外，本文还提供了一个专用的利用SRAM的Softmax内核和一个作为基线的CPU实现。在Grayskull上，Softmax操作消耗了大部分注意力权重计算的运行时间。专用的Softmax内核与CPU实现相比，速度提高了高达10倍，而融合内核中的Softmax实现比专用的Softmax内核快大约1.8倍。所有实现的时间和内存复杂度都是序列长度的平方。目前，Grayskull e150对于一般公众来说大约比Nvidia H100 PCIe（一款最先进的GPU）便宜30倍，提供的SRAM大约是其1.5倍。

本论文旨在解决Transformer的self-attention层在使用SRAM而非DRAM时的速度问题，提出了一种融合内核的方案。

论文提出了一种融合内核的方案，该方案将矩阵乘法、注意力得分缩放和Softmax操作结合在一起，专门利用Grayskull的大型SRAM。此外，还提出了一种专用的Softmax内核和一个基于CPU的实现作为基准。该方案能够显著提高Softmax操作的运行速度。

论文的亮点包括：1. 提出了一种新的融合内核的方案，能够显著提高Softmax操作的运行速度；2. 专门利用Grayskull的大型SRAM，相比当前的GPU方案更加经济实惠；3. 实验结果表明，该方案的性能比CPU实现快$10 imes$，比专用的Softmax内核快$1.8 imes$。

在这个领域中，最近的相关研究包括：1. Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context；2. Reformer: The Efficient Transformer；3. Longformer: The Long-Document Transformer。