SparkAttention: High-Performance Multi-Head Attention for Large Models on Volta GPU Architecture
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2025
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866913998906064896 |
|---|---|
| author | Xu, Youxuan Wu, Tong Li, Shigang Wang, Xueying Wang, Jingjing |
| author_facet | Xu, Youxuan Wu, Tong Li, Shigang Wang, Xueying Wang, Jingjing |
| contents | Transformer are widely used in various fields such as natural language processing and computer vision. However, the training time for large Transformer models can be challenging due to the Multi-Head Attention (MHA) mechanism. Especially as models become larger, training becomes more costly. So it is crucial to utilize various resources for efficient model training. Currently, NVIDIA Volta GPU is still widely used. However, because the computational shapes supported by Tensor Core Units (TCU) of Volta GPU differ from other GPU architectures, most efforts have not focused on using them to accelerate Transformer training. To address this issue, we propose SparkAttention, an acceleration library designed to speed up MHA training on the Volta GPU. SparkAttention leverages TCU and kernel fusion to reduce the number of high bandwidth memory (HBM) accesses and overhead. Our End-to-End experimental results on an NVIDIA V100 GPU show that SparkAttention achieves on average 1.80$\times$ (up to 2.46$\times$) speedup compared to using PyTorch. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_12784 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | SparkAttention: High-Performance Multi-Head Attention for Large Models on Volta GPU Architecture Xu, Youxuan Wu, Tong Li, Shigang Wang, Xueying Wang, Jingjing Distributed, Parallel, and Cluster Computing C.1.4; I.2.11 Transformer are widely used in various fields such as natural language processing and computer vision. However, the training time for large Transformer models can be challenging due to the Multi-Head Attention (MHA) mechanism. Especially as models become larger, training becomes more costly. So it is crucial to utilize various resources for efficient model training. Currently, NVIDIA Volta GPU is still widely used. However, because the computational shapes supported by Tensor Core Units (TCU) of Volta GPU differ from other GPU architectures, most efforts have not focused on using them to accelerate Transformer training. To address this issue, we propose SparkAttention, an acceleration library designed to speed up MHA training on the Volta GPU. SparkAttention leverages TCU and kernel fusion to reduce the number of high bandwidth memory (HBM) accesses and overhead. Our End-to-End experimental results on an NVIDIA V100 GPU show that SparkAttention achieves on average 1.80$\times$ (up to 2.46$\times$) speedup compared to using PyTorch. |
| title | SparkAttention: High-Performance Multi-Head Attention for Large Models on Volta GPU Architecture |
| topic | Distributed, Parallel, and Cluster Computing C.1.4; I.2.11 |
| url | https://arxiv.org/abs/2502.12784 |