flagos-ai/FlagGems

何を解決するか

FlagGems は、AIチップセット固有のソフトウェアスタック間の断片化を解消するように設計されています。異なるAIアクセラレータ用に別々のコードベースを書き直すか維持する必要がなくなるため、LLMのトレーニングおよび推論において「一度開発すれば、どこでも実行可能」なワークフローを実現します。

仕組み

Triton言語で実装された高性能で汎用的なオペレータライブラリです。PyTorch ATenバックエンドに登録することで、開発者はなじみ深いPyTorch APIを使用しつつ、下位の演算がTritonカーネルによって実行されるようになります。これにより、10種類以上のハードウェアプラットフォームをサポートするバックエンド非依存のアプローチが可能になり、ユーザーが低レベルAPIを変更する必要がありません。

対象ユーザー

• PyTorchコードを変更せずにハードウェアアクセラレーションを活用したいモデル開発者。
• CUDAよりも読みやすく使いやすい言語で高性能カーネルを書きたいカーネル開発者。

特徴

• 広範なハードウェア対応: 多様なAIアクセラレータ向けに10種類以上のバックエンドをサポート。
• PyTorch互換性: PyTorch APIとシームレスに統合され、torch.compileに依存せず、イージーモードでも利用可能。
• Tritonベース: CUDAと同等のパフォーマンスを実現しつつ、より読みやすいコードを提供。
• 自動コード生成: 任意の入力型やレイアウトをサポートする自動点対点オペレータコード生成を含む。
• 効率的なディスパッチ: 関数ごとの実行時カーネルディスパッチが高速。