Attention in SRAM on Tenstorrent Grayskull

Abstract: 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.

• The paper introduces a fused kernel that integrates matrix multiplication, attention scaling, and Softmax using on-chip SRAM to accelerate self-attention.
• It demonstrates up to a 10-fold speedup over CPU implementations while reducing memory access and kernel dispatch overhead.
• The study highlights cost-effective AI acceleration and proposes further integration of query, key, and value computations for enhanced performance.

Attention in SRAM on Tenstorrent Grayskull

The paper "Attention in SRAM on Tenstorrent Grayskull" by Moritz Thüning rigorously examines the potential performance enhancements achievable by employing SRAM to execute the Transformer's self-attention mechanism, a pivotal component in contemporary machine learning architectures, particularly in NLP. This investigation revolves around the Tenstorrent Grayskull architecture, a hardware platform distinguished by a large, distributed SRAM grid designed to optimize AI computational tasks.

Key Contributions

The core contribution of this work is the development and evaluation of a fused kernel for the Grayskull e150 architecture. This kernel leverages the on-chip SRAM to perform essential operations in the self-attention layer including matrix multiplication, attention score scaling, and the Softmax function. The study presents a specialized Softmax kernel as well as a baseline CPU implementation for comparative analysis.

Numerical Results

The paper provides compelling empirical data demonstrating significant performance improvements with the proposed methodologies:

• The dedicated Softmax kernel on Grayskull achieves up to a 10-fold speedup over the baseline CPU implementation.
• The Softmax component within the fused kernel operates approximately 1.8 times faster than the standalone Softmax kernel.
• Although the enhancements from the dedicated Softmax to the fused kernel are not linear, the fused kernel generally outperforms individual operations by reducing memory access and dispatch overhead.

The analysis also notes that the time and memory complexities associated with these implementations scale quadratically with sequence length, indicative of inherent challenges in scaling attention mechanisms.

Practical and Theoretical Implications

Practical Implications

From a practical standpoint, these advancements indicate substantial cost and performance benefits:

• The Grayskull e150 card is highlighted as being approximately 30 times more affordable than an Nvidia H100 PCIe card, making it an attractive option for cost-sensitive applications.
• The significant speedup in attention computations can enhance real-time processing capabilities in NLP applications, robotics, and other AI domains where rapid inference is critical.

Theoretical Implications

The research accentuates the capability of custom hardware architectures to mitigate limitations of existing AI accelerators, particularly concerning memory bandwidth and latency constraints. The critical examination of SRAM utilization provides a roadmap for future architectural optimizations.

Runtime Observations

Comprehensive runtime analysis underscores that the Softmax operation is a significant bottleneck in self-attention layers, a finding corroborated by previous works such as FlashAttention. The fused kernel demonstrates improved efficiency by integrating matrix multiplication and Softmax operations, thereby capitalizing on data proximity within the SRAM.

Future Directions

The paper proposes several avenues for further exploration:

1. Extended Integration: Future efforts could aim to incorporate the computation of queries, keys, and values directly into the fused kernel to further minimize external memory access and kernel dispatching overheads.
2. Enhanced Softmax Variants: Investigating alternative Softmax computation strategies, such as centralized calculation of global sums and maxima, could provide additional energy and efficiency gains.
3. Porting and Scaling: Extension to newer hardware generations like Tenstorrent Wormhole and scaling across multiple Grayskull cards could unlock further performance improvements and insights into multi-card operation dynamics.

Conclusion

The study "Attention in SRAM on Tenstorrent Grayskull" effectively demonstrates the potential of specialized hardware architectures in enhancing the efficiency of self-attention mechanisms within the Transformer model. The quantitative results combined with the practical and theoretical implications underscore the utility of SRAM in AI accelerators and set the stage for future innovations in efficient machine learning hardware design.