openxla/tokamax

何を解決するか

Tokamax は NVIDIA GPU および Google TPU 用の高性能カスタムアクセラレータカーネルを提供し、一般的なディープラーニング演算の低レベル最適化を手動で実装する必要を減らします。特に、現代の AI モデルに不可欠な Attention や正規化などの演算における性能ギャップに焦点を当てています。

仕組み

JAX と Pallas の上に構築された Tokamax は、特定の演算向けの最先端のカーネルを実装しています。ユーザーは特定の実装（例：Triton や Mosaic）を選択できるか、カーネルの形状に基づいて最も効率的なものを自動的に選択させることができます。パフォーマンスを最大化するために、特定のワークロードに最適な設定を探索する自動チューニングシステムを備えており、その設定はシリアライズして再利用可能で、数値的一貫性を保証します。

対象ユーザー

JAX を使用し、カスタムカーネルと自動チューニングを通じて GPU および TPU ハードウェアの最大性能を引き出したい AI 研究者およびエンジニア向けに設計されています。

主な特徴

• マルチアクセラレータ対応：NVIDIA GPU と Google TPU の両方で動作します。
• 事前構築されたカーネル：FlashAttention、ゲート付き線形ユニット（SwiGLU）、LayerNorm、RMSNorm、Mixture of Experts（ragged_dot）の最適化実装を含みます。
• 自動チューニングツールキット：最速のカーネル設定を自動的に見つけるフレームワークを提供し、ユーザーが独自の自動チューニング探索空間を定義できるようにしています。
• ハードウェアベンチマーク：JAX の Python オーバーヘッドを回避し、実際のアクセラレータ実行時間を測定するユーティリティを備えており、CUPTI プロファイラなどのツールに対応しています。