Accelerating Sparse Transformer Inference on GPU
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2025
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866916056714444800 |
|---|---|
| author | Dai, Wenhao Deng, Haodong Rong, Mengfei Yang, Xinyu Liu, Hongyu Liu, Fangxin Yang, Hailong Cao, Qianwen Sun, Qingxiao |
| author_facet | Dai, Wenhao Deng, Haodong Rong, Mengfei Yang, Xinyu Liu, Hongyu Liu, Fangxin Yang, Hailong Cao, Qianwen Sun, Qingxiao |
| contents | Large language models (LLMs) are popular around the world due to their powerful understanding capabilities. As the core component of LLMs, accelerating Transformer through parallelization has gradually become a hot research topic. Mask layers introduce sparsity into Transformer to reduce calculations. However, previous works rarely focus on the performance optimization of sparse Transformer. In addition, current static operator fusion schemes fail to adapt to diverse application scenarios. To address the above problems, we propose STOF, a framework that incorporates optimizations for Sparse Transformer that enables flexible masking and Operator Fusion on GPU. For multi-head attention (MHA) structure, STOF maps the computation to row-wise or blockwise kernels with unique storage formats according to analytical modeling. For downstream operators, STOF maps the fusion scheme to compilation templates and determines the optimal running configuration through two-stage searching. The experimental results show that compared to the stateof-the-art work, STOF achieves maximum speedups of 1.6x in MHA computation and 1.4x in end-to-end inference. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_06095 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Accelerating Sparse Transformer Inference on GPU Dai, Wenhao Deng, Haodong Rong, Mengfei Yang, Xinyu Liu, Hongyu Liu, Fangxin Yang, Hailong Cao, Qianwen Sun, Qingxiao Machine Learning Large language models (LLMs) are popular around the world due to their powerful understanding capabilities. As the core component of LLMs, accelerating Transformer through parallelization has gradually become a hot research topic. Mask layers introduce sparsity into Transformer to reduce calculations. However, previous works rarely focus on the performance optimization of sparse Transformer. In addition, current static operator fusion schemes fail to adapt to diverse application scenarios. To address the above problems, we propose STOF, a framework that incorporates optimizations for Sparse Transformer that enables flexible masking and Operator Fusion on GPU. For multi-head attention (MHA) structure, STOF maps the computation to row-wise or blockwise kernels with unique storage formats according to analytical modeling. For downstream operators, STOF maps the fusion scheme to compilation templates and determines the optimal running configuration through two-stage searching. The experimental results show that compared to the stateof-the-art work, STOF achieves maximum speedups of 1.6x in MHA computation and 1.4x in end-to-end inference. |
| title | Accelerating Sparse Transformer Inference on GPU |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2506.06095 |