about
Attention in SRAM on Tenstorrent Grayskull (1.5x SRAM, 30x cheaper than H100) (arxiv.org)
3 points by molli on Jul 23, 2024 | hide | past | pdf | 1 comment on HN

In plain words: A new routine keeps a key attention step in AI language models—matrix multiply, score scaling, and Softmax—inside a chip's fast on-chip memory instead of slower main memory. Its Softmax runs about 1.8 times faster than the same operation written as a separate routine.

Abstract · Attention in SRAM on Tenstorrent Grayskull

When implementations of the Transformer's self-attention layer utilize SRAM instead of DRAM, they can achieve significant speedups. The Tenstorrent Grayskull architecture provides a large SRAM, distributed across a grid of cores. This work presents a fused kernel for Grayskull, that exclusively utilizes its large SRAM by combining matrix multiplication, attention score scaling and Softmax operations. Additionally, a dedicated Softmax kernel utilizing the SRAM and a CPU implementation serving as a baseline are presented. The Softmax operation consumes most of the runtime in the computation of attention weights from queries and keys on Grayskull. The speedup of the dedicated Softmax kernel compared to the CPU implementation is up to $10 \times$, and the Softmax implementation inside the fused kernel is approximately $1.8 \times$ faster than the dedicated Softmax kernel. The time and memory complexity of all implementations is quadratic in sequence length. Currently, the Grayskull e150 is approximately $30 \times$ cheaper for the general public than an Nvidia H100 PCIe (a state-of-the-art GPU) and offers approximately $1.5 \times$ more SRAM.

Moritz Thüning
arXiv:2407.13885 · cs.LG, cs.PF · submitted Jul 18, 2024
abstract · pdf · html · 8 pages, 6 figures, code: https://github.com/moritztng/grayskull-attention

add comment on HN
Also discussed: Aug 2024 (1 point, 0 comments)

When implementations of the Transformer's self-attention layer utilize SRAM instead of DRAM, they can achieve significant speedups. The Tenstorrent Grayskull architecture provides a large SRAM, distributed across a grid of cores. This work presents a fused kernel for Grayskull, that exclusively utilizes its large SRAM by combining matrix multiplication, attention score scaling and Softmax operations. Additionally, a dedicated Softmax kernel utilizing the SRAM and a CPU implementation serving as a baseline are presented. The Softmax operation consumes most of the runtime in the computation of attention weights from queries and keys on Grayskull. The speedup of the dedicated Softmax kernel compared to the CPU implementation is up to 10x, and the Softmax implementation inside the fused kernel is approximately 1.8x faster than the dedicated Softmax kernel. The time and memory complexity of all implementations is quadratic in sequence length. Currently, the Grayskull e150 is approximately 30x cheaper for the general public than an Nvidia H100 PCIe (a state-of-the-art GPU) and offers approximately 1.5x more SRAM.

Code: https://github.com/moritztng/grayskull-attention