openxla/tokamax

解决的问题

Tokamax 为 NVIDIA GPU 和 Google TPU 提供高性能的自定义加速器内核，减少了开发者手动实现常见深度学习操作低级优化的需求。它特别针对现代 AI 模型中至关重要的注意力（attention）和归一化等操作中存在的性能差距。

工作原理

基于 JAX 和 Pallas 构建，Tokamax 实现了针对特定操作的前沿内核。用户可以选择特定实现（例如 Triton 或 Mosaic），或让库根据内核形状自动选择最高效的实现。为最大化性能，它包含一个自动调优系统，可为特定工作负载搜索最佳配置，该配置可序列化并重复使用，以确保数值一致性。

适用人群

专为使用 JAX 的 AI 研究人员和工程师设计，他们希望通过自定义内核和自动调优从 GPU 和 TPU 硬件中榨取最大性能。

主要亮点

• 多加速器支持：支持 NVIDIA GPU 和 Google TPU 两种平台。
• 预构建内核：包含 FlashAttention、门控线性单元（SwiGLU）、LayerNorm、RMSNorm 以及混合专家（ragged_dot）的优化实现。
• 自动调优工具：提供框架以自动寻找最快内核配置，并允许用户自定义自动调优的搜索空间。
• 硬件基准测试：包含测量实际加速器执行时间的工具，绕过 JAX 的 Python 开销，支持 CUPTI 等分析工具。