INT-FlashAttention: Enabling Flash Attention for INT8 Quantization
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866910620760145920 |
|---|---|
| author | Chen, Shimao Liu, Zirui Wu, Zhiying Zheng, Ce Cong, Peizhuang Jiang, Zihan Wu, Yuhan Su, Lei Yang, Tong |
| author_facet | Chen, Shimao Liu, Zirui Wu, Zhiying Zheng, Ce Cong, Peizhuang Jiang, Zihan Wu, Yuhan Su, Lei Yang, Tong |
| contents | As the foundation of large language models (LLMs), self-attention module faces the challenge of quadratic time and memory complexity with respect to sequence length. FlashAttention accelerates attention computation and reduces its memory usage by leveraging the GPU memory hierarchy. A promising research direction is to integrate FlashAttention with quantization methods. This paper introduces INT-FlashAttention, the first INT8 quantization architecture compatible with the forward workflow of FlashAttention, which significantly improves the inference speed of FlashAttention on Ampere GPUs. We implement our INT-FlashAttention prototype with fully INT8 activations and general matrix-multiplication (GEMM) kernels, making it the first attention operator with fully INT8 input. As a general token-level post-training quantization framework, INT-FlashAttention is also compatible with other data formats like INT4, etc. Experimental results show INT-FlashAttention achieves 72% faster inference speed and 82% smaller quantization error compared to standard FlashAttention with FP16 and FP8 data format. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2409_16997 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | INT-FlashAttention: Enabling Flash Attention for INT8 Quantization Chen, Shimao Liu, Zirui Wu, Zhiying Zheng, Ce Cong, Peizhuang Jiang, Zihan Wu, Yuhan Su, Lei Yang, Tong Machine Learning Artificial Intelligence As the foundation of large language models (LLMs), self-attention module faces the challenge of quadratic time and memory complexity with respect to sequence length. FlashAttention accelerates attention computation and reduces its memory usage by leveraging the GPU memory hierarchy. A promising research direction is to integrate FlashAttention with quantization methods. This paper introduces INT-FlashAttention, the first INT8 quantization architecture compatible with the forward workflow of FlashAttention, which significantly improves the inference speed of FlashAttention on Ampere GPUs. We implement our INT-FlashAttention prototype with fully INT8 activations and general matrix-multiplication (GEMM) kernels, making it the first attention operator with fully INT8 input. As a general token-level post-training quantization framework, INT-FlashAttention is also compatible with other data formats like INT4, etc. Experimental results show INT-FlashAttention achieves 72% faster inference speed and 82% smaller quantization error compared to standard FlashAttention with FP16 and FP8 data format. |
| title | INT-FlashAttention: Enabling Flash Attention for INT8 Quantization |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2409.16997 |