QiMeng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866915343453126656 |
|---|---|
| author | Zhou, Qirui Peng, Shaohui Xiong, Weiqiang Chen, Haixin Wen, Yuanbo Li, Haochen Li, Ling Guo, Qi Zhao, Yongwei Gao, Ke Chen, Ruizhi Wu, Yanjun Zhao, Chen Chen, Yunji |
| author_facet | Zhou, Qirui Peng, Shaohui Xiong, Weiqiang Chen, Haixin Wen, Yuanbo Li, Haochen Li, Ling Guo, Qi Zhao, Yongwei Gao, Ke Chen, Ruizhi Wu, Yanjun Zhao, Chen Chen, Yunji |
| contents | The attention operator remains a critical performance bottleneck in large language models (LLMs), particularly for long-context scenarios. While FlashAttention is the most widely used and effective GPU-aware acceleration algorithm, it must require time-consuming and hardware-specific manual implementation, limiting adaptability across GPU architectures. Existing LLMs have shown a lot of promise in code generation tasks, but struggle to generate high-performance attention code. The key challenge is it cannot comprehend the complex data flow and computation process of the attention operator and utilize low-level primitive to exploit GPU performance.
To address the above challenge, we propose an LLM-friendly Thinking Language (LLM-TL) to help LLMs decouple the generation of high-level optimization logic and low-level implementation on GPU, and enhance LLMs' understanding of attention operator. Along with a 2-stage reasoning workflow, TL-Code generation and translation, the LLMs can automatically generate FlashAttention implementation on diverse GPUs, establishing a self-optimizing paradigm for generating high-performance attention operators in attention-centric algorithms. Verified on A100, RTX8000, and T4 GPUs, the performance of our methods significantly outshines that of vanilla LLMs, achieving a speed-up of up to 35.16x. Besides, our method not only surpasses human-optimized libraries (cuDNN and official library) in most scenarios but also extends support to unsupported hardware and data types, reducing development time from months to minutes compared with human experts. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_12355 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | QiMeng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm Zhou, Qirui Peng, Shaohui Xiong, Weiqiang Chen, Haixin Wen, Yuanbo Li, Haochen Li, Ling Guo, Qi Zhao, Yongwei Gao, Ke Chen, Ruizhi Wu, Yanjun Zhao, Chen Chen, Yunji Machine Learning Computation and Language I.2.7 The attention operator remains a critical performance bottleneck in large language models (LLMs), particularly for long-context scenarios. While FlashAttention is the most widely used and effective GPU-aware acceleration algorithm, it must require time-consuming and hardware-specific manual implementation, limiting adaptability across GPU architectures. Existing LLMs have shown a lot of promise in code generation tasks, but struggle to generate high-performance attention code. The key challenge is it cannot comprehend the complex data flow and computation process of the attention operator and utilize low-level primitive to exploit GPU performance. To address the above challenge, we propose an LLM-friendly Thinking Language (LLM-TL) to help LLMs decouple the generation of high-level optimization logic and low-level implementation on GPU, and enhance LLMs' understanding of attention operator. Along with a 2-stage reasoning workflow, TL-Code generation and translation, the LLMs can automatically generate FlashAttention implementation on diverse GPUs, establishing a self-optimizing paradigm for generating high-performance attention operators in attention-centric algorithms. Verified on A100, RTX8000, and T4 GPUs, the performance of our methods significantly outshines that of vanilla LLMs, achieving a speed-up of up to 35.16x. Besides, our method not only surpasses human-optimized libraries (cuDNN and official library) in most scenarios but also extends support to unsupported hardware and data types, reducing development time from months to minutes compared with human experts. |
| title | QiMeng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm |
| topic | Machine Learning Computation and Language I.2.7 |
| url | https://arxiv.org/abs/2506.12355 |