Tenstorrent Grayskull 上 SRAM 中的注意力机制

当 Transformer 的自注意力层实现使用 SRAM 而非 DRAM 时，可以实现显著的加速。Tenstorrent Grayskull 架构提供了跨越网格核心分布的大量 SRAM。本工作为 Grayskull 提出一种融合内核，该内核完全利用其大量 SRAM，通过组合矩阵乘法、注意力得分缩放和 Softmax 操作。此外，还提出一个专用的 Softmax 内核，利用 SRAM 作为基准。Softmax 操作在计算注意力权重（从查询和键）时占据大部分运行时间。相对于 CPU 实现，专用 Softmax 内核的加速比最高可达 $10 \times$ ，而融合内核中 Softmax 的实现大约比专用 Softmax 内核快 $1.8 \times$ 。所有实现的时间和空间复杂度均随序列长度呈二次增长。当前，Grayskull e150 相较于 Nvidia H100 PCIe（一种最先进的 GPU）对一般大众的价格便宜约 $30 \times$ ，且提供约 $1.5 \times$ 的 SRAM。