250617_QiMeng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm

论文信息

  • 标题: QiMeng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm

  • arXiv ID: 2506.12355

  • 作者: Authors:Qirui Zhou, Shaohui Peng, Weiqiang Xiong, Haixin Chen, Yuanbo Wen, Haochen Li, Ling Li, Qi Guo, Yongwei Zhao, Ke Gao, Ruizhi Chen, Yanjun Wu, Chen Zhao, Yunji Chen

  • 发表日期: 2025-06-14T05:38:19+00:00

  • 论文链接: 2506.12355

  • 总结生成时间: 2025-06-17 19:26:49

一句话概要
QiMeng-Attention 通过设计一种面向大语言模型的思维语言（LLM-TL）和两阶段推理流程，实现了跨 GPU 架构的高性能注意力算子自动生成，解决了现有手动优化方法耗时且硬件依赖性强的问题。

主体
注意力算子在处理长上下文的大语言模型中已成为关键性能瓶颈，而当前主流的 FlashAttention 加速算法需要针对不同 GPU 架构进行耗时的手动优化，严重限制了其通用性。尽管大语言模型在代码生成任务中展现出潜力，但直接生成高性能注意力算子代码仍面临挑战，主要源于模型难以理解复杂的数据流、计算过程以及有效利用底层 GPU 原语。

为解决这一问题，作者提出了一种创新的思维语言 LLM-TL，通过将高层优化逻辑与底层 GPU 实现解耦，显著提升了大语言模型对注意力算子的理解能力。配合两阶段推理流程（TL 代码生成与翻译），该方法能够自动为不同 GPU 生成 FlashAttention 实现，形成了一种自优化的高性能注意力算子生成范式。这种设计不仅降低了硬件适配的复杂度，还使模型能够自主探索最优计算路径。

在 A100、RTX8000 和 T4 等 GPU 上的实验表明，该方法生成的算子性能远超原始大语言模型生成结果，最高加速比达 35.16 倍。更值得注意的是，其表现甚至优于人工优化的 cuDNN 等官方库，同时支持了原有库未覆盖的硬件和数据类型，将开发周期从数月缩短至分钟级。这一突破为注意力机制的硬件适配提供了标准化解决方案，同时验证了大语言模型在复杂系统优化中的潜力。

最后一句
该研究为自动生成高性能计算内核开辟了新路径，未来或可扩展至其他计算密集型算子的跨平台优化领域。