用Triton编写高效跨GPU的注意力核，实现媲美甚至超越手工优化的性能。
The Anatomy of a Triton Attention Kernel
- 纯Triton语言实现可移植的分页注意力核，无需硬件定制。
- 在NVIDIA和AMD GPU上均达最优性能，效率提升超4倍（19.7%→105.9%）。
- 适合关注模型跨平台部署与开源编译器优化的研究者与工程师。
产业界与学术界长期追求一种可在多种硬件架构间移植、无需低层手工调优且仍具备顶尖效率的LLM推理平台。本文证明此类平台确实可行，并分享实践经验。我们开发了一种前沿的分页注意力核，作为众多LLM部署中的核心性能组件，完全基于领域专用即时编译语言Triton构建，在NVIDIA与AMD GPU上均实现顶尖性能。文中详述了整体设计思路、关键算法与系统级改进、用于解锁效率的参数自动调优方法，以及集成至主流推理服务器所必需的操作。通过这些工作，通用Triton注意力核的性能从原先的19.7%提升至105.9%，显著超越现有最佳方案。结果表明，开源领域专用语言能有效推动模型在不同GPU厂商间的可移植性。
原文摘要 · Abstract (English)
A long-standing goal in both industry and academia is to develop an LLM inference platform that is portable across hardware architectures, eliminates the need for low-level hand-tuning, and still delivers best-in-class efficiency. In this work, we demonstrate that portable, efficient cross-platform LLM inference is indeed possible and share our experience. We develop a state-of-the-art paged attention kernel, the core performance-critical component of many LLM deployments, that builds exclusively on the domain-specific just-in-time compiled language Triton to achieve state-of-the-art performance on both NVIDIA and AMD GPUs. We describe our high-level approach, the key algorithmic and system-level improvements, the parameter auto-tuning required to unlock efficiency, and the integrations into a popular inference server that are necessary to bring the performance of a generic Triton attention kernel from 19.7% of the state-of-the-art to 105.9%. Our results highlight how open-source domain-specific languages can be leveraged to unlock model portability across different GPU vendors.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。