ROCm/aiter

解決的問題

AITER（AI Tensor Engine for ROCm）為AMD GPU提供高性能的優化GPU核心庫。它解決了框架開發者在無需從零開始撰寫低階GPU核心的情況下，獲得可用於生產環境、高度優化的AI推論與訓練運算子的需求。

工作原理

它作為一個統一的運算子集合，可整合至各種AI框架中。AITER利用多種核心後端以實現最大效能，包括Triton、可組合核心（CK）與手動調校的組合語言（ASM）。它也利用FlyDSL實現特定運算子（如GEMM和MoE），並包含一個輕量級C++模板庫Opus，以加速HIP核心的開發與建置時間。

適用對象

此庫主要針對框架開發者（如vLLM、SGLang或JAX的開發者）以及需要在AMD Instinct與Radeon GPU上優化工作負載的自訂AI推論引擎工程師。

主要亮點

• 廣泛的框架支援：作為ROCm上LLM推論的預設核心後端，已整合至vLLM、SGLang與ATOM中。
• 多樣化的運算子集合：包含針對注意力機制（MHA、MLA、分頁注意力）、專家混合（MoE）、GEMM、歸一化與量化等的優化核心。
• 多後端方法：結合Triton、CK與手動調校的ASM，針對不同硬體架構進行優化。
• 顯著的效能提升：展現出顯著的加速效果，例如MLA解碼核心最高達17倍，MHA預填入核心最高達14倍。
• 彈性的API：提供C++與Python API，便於整合。