# Attention in SRAM on Tenstorrent Grayskull

arXiv: https://arxiv.org/abs/2407.13885 · HTML: https://mlsanity.com/paper/2407.13885/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 19 references that are in the map.
- First posted on arXiv 18 Jul 2024.

## Builds on (most cited first, 9 of 9)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Efficient Content-Based Sparse Attention with Routing Transformers](https://mlsanity.com/paper/2003.05997/) (2020)
- [Scatterbrain: Unifying Sparse and Low-rank Attention Approximation](https://mlsanity.com/paper/2110.15343/) (2021)

## Built on by (most cited first, 1 of 1)

- [Operator Fusion for LLM Inference on the Tensix Architecture](https://mlsanity.com/paper/2606.09879/) (2026)

## Abstract (arXiv)

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 ×, and the Softmax implementation inside the fused kernel is approximately 1.8 × 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 × cheaper for the general public than an Nvidia H100 PCIe (a state-of-the-art GPU) and offers approximately 1.5 × more SRAM.