arXiv:2506.12355cs.LGcs.CL2025-06ACL被引 5

用大模型自动生成高性能注意力代码,支持多种显卡且速度提升35倍。

QiMeng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm

  • 设计新型思维语言,让大模型分阶段生成并转换高性能注意力代码。
  • 在A100、RTX8000等多款显卡上实现最高35.16倍加速。
  • 自动适配新硬件与数据类型,开发时间从月级缩短至分钟级。

注意力算子仍是大语言模型(LLM)在长序列场景下的关键性能瓶颈。尽管FlashAttention是目前最广泛使用且高效的GPU感知加速算法,但其需耗时且依赖硬件的手工实现,限制了跨GPU架构的适应性。现有LLM在代码生成任务中虽有潜力,却难以生成高性能注意力代码,核心难题在于无法理解注意力算子复杂的数据流与计算过程,并有效利用底层原语挖掘GPU性能。为此,我们提出一种面向大模型的思维语言(LLM-TL),帮助大模型分离高层次优化逻辑与低层次GPU实现,增强对注意力算子的理解。结合两阶段推理流程——代码生成与翻译,大模型可自动在多样GPU上生成FlashAttention实现,建立注意力为中心算法的自优化生成范式。在A100、RTX8000和T4 GPU上的验证表明,本方法性能显著优于原始大模型,最快提速达35.16倍。此外,该方法在多数场景下超越人工优化库(cuDNN与官方库),并扩展支持未受支持的硬件与数据类型,开发时间由数月缩短至数分钟。

原文摘要 · Abstract (English)

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.

注意力机制代码生成自优化GPU加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。