01 / UNDERSTAND
仕組みと使いどころ
NKIは、AWSのAIチップが備えるハードウェアプリミティブや命令へ低レベルにアクセスし、独自機能や最適化を組み込んだ計算カーネルを開発するための環境です。Tritonに似た構文とタイル単位のセマンティクスを採用しています。NKI 0.3.0では安定版となり、APIなどのソースを公開した標準ライブラリと、通常のCPU上でカーネルを記述・テスト・デバッグできるCPU Simulatorが提供されています。NKI自体は汎用的な機械学習フレームワークではなく、作成したカーネルの性能向上を一律に保証するものでもありません。
02 / IN PRACTICE
Neuron Kernel Interfaceの活用例
たとえば、こんな場面で
モデル内の重要な演算に既存の実装では足りない機能や最適化が必要な場合、NKIでTrainiumまたはInferentia向けのカスタムカーネルを試作する、といった活用が考えられます。CPU Simulatorを使えば、対応ハードウェアがない環境でも記述やデバッグを進められます。
03 / GOOD TO KNOW
使うときのポイント
利用時は、対象のTrainiumまたはInferentia世代、NKIとNeuron SDKのバージョン、使用するAPIやカーネルが安定版か実験的機能かを公式資料で確認してください。
参考資料
AWS Trainium Research – Amazon Web Services （新しいタブで開きます）AWS Neuron introduces Neuron Kernel Interface (NKI), NxD Training, and JAX support for training - AWS （新しいタブで開きます）AWS Neuron SDK 2.29.0 now available with NKI and Neuron Explorer out of Beta, a new CPU Simulator, and an expanded NKI Library - AWS （新しいタブで開きます）用語の説明や活用例に関連する参考資料です。
READ & LEARN
記事一覧を見る 用語を知ったら、活用を知る。
AI LABで、AIの動向や業務に取り入れるヒントを読む。