AttentionEngine让注意力机制在不同硬件上自动高效运行，性能提升最高达10倍。
AttentionEngine: A Versatile Framework for Efficient Attention Mechanisms on Diverse Hardware Platforms
- 将注意力计算拆解为可定制模块，灵活适配多种算法需求。
- 通过模板化代码与跨平台调度，实现无需人工干预的性能优化。
- 适合需要在多设备部署大模型的研究者和开发者使用。
Transformer 和大型语言模型（LLMs）彻底改变了机器学习，其成功核心在于注意力机制。随着注意力变体不断增多，跨异构硬件平台的性能优化面临更大挑战。现有优化策略通常过于局限，需大量手动调整以适应模型配置或硬件环境变化。本文提出 AttentionEngine，一个全面的框架，旨在统一优化多样硬件后端上的注意力机制。该框架将注意力计算分解为模块化操作，并支持可定制组件，从而灵活适配不同算法需求。同时，通过可编程模板与稳健的跨平台调度策略，实现内核优化的自动化。实验证明，在现有方法无法覆盖的配置下，性能提升最高达10倍。AttentionEngine为开发和部署注意力机制提供了可扩展、高效的基础设施，显著减少人工调优。代码已开源，地址为 https://github.com/microsoft/AttentionEngine。
原文摘要 · Abstract (English)
Transformers and large language models (LLMs) have revolutionized machine learning, with attention mechanisms at the core of their success. As the landscape of attention variants expands, so too do the challenges of optimizing their performance, particularly across different hardware platforms. Current optimization strategies are often narrowly focused, requiring extensive manual intervention to accommodate changes in model configurations or hardware environments. In this paper, we introduce AttentionEngine, a comprehensive framework designed to streamline the optimization of attention mechanisms across heterogeneous hardware backends. By decomposing attention computation into modular operations with customizable components, AttentionEngine enables flexible adaptation to diverse algorithmic requirements. The framework further automates kernel optimization through a combination of programmable templates and a robust cross-platform scheduling strategy. Empirical results reveal performance gains of up to 10x on configurations beyond the reach of existing methods. AttentionEngine offers a scalable, efficient foundation for developing and deploying attention mechanisms with minimal manual tuning. Our code has been open-sourced and is available at https://github.com/microsoft/AttentionEngine.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。