QiMeng-Attention：当前最优注意力算子由当前最优注意力算法生成
机器学习
2025-06-17 v1 计算与语言
摘要
注意力算子在大语言模型（LLMs）中仍然是一个关键的性能瓶颈，尤其是在长上下文场景中。虽然 FlashAttention 是最广泛使用且最有效的 GPU 感知加速算法，但它需要耗时且依赖硬件的手动实现，限制了在不同 GPU 架构上的适应性。现有的大语言模型在代码生成任务中展现出了很大潜力，但在生成高性能注意力代码方面仍存在困难。关键挑战在于它无法理解注意力算子中复杂的数据流和计算过程，也无法利用底层原语来挖掘 GPU 性能。为了解决上述挑战，我们提出了一种面向大语言模型的思考语言（LLM-TL），帮助大语言模型将高层优化逻辑的生成与 GPU 底层实现解耦，并增强大语言模型对注意力算子的理解。结合两阶段推理工作流——TL-Code 生成与翻译，大语言模型可以自动在多种 GPU 上生成 FlashAttention 实现，建立用于生成注意力中心算法中高性能注意力算子的自优化范式。在 A100、RTX8000 和 T4 GPU 上验证后，我们方法的性能显著优于普通大语言模型，实现了最高 35.16 倍的加速。此外，我们的方法在大多数场景下不仅超越了人工优化的库（cuDNN 和官方库），还扩展了对不受支持硬件和数据类型的支持，将开发时间从数月缩短至数分钟，相比人工专家大幅减少。
引用
@article{arxiv.2506.12355,
title = {QiMeng-Attention: SOTA Attention Operator is generated by SOTA Attention Algorithm},
author = {Qirui Zhou and Shaohui Peng and Weiqiang Xiong and Haixin Chen and Yuanbo Wen and Haochen Li and Ling Li and Qi Guo and Yongwei Zhao and Ke Gao and Ruizhi Chen and Yanjun Wu and Chen Zhao and Yunji Chen},
journal= {arXiv preprint arXiv:2506.12355},
year = {2025}
}