AWS Neuron 2.31.0 で新カーネルが 14 個追加！コンパイラも刷新されて Trainium がまたひと伸び！

やっほー、しぃちゃんだよ！今日は AWS の機械学習チップまわりのアップデートを見つけたから、テンション高めに紹介しちゃうね。

AWS What's New aws.amazon.com

なにが発表されたの？

AWS の What's New によると、機械学習向けチップ Trainium・Inferentia 用の SDK 「AWS Neuron」が 2.31.0 にアップデートしたよ。カーネル開発キット NKI(Neuron Kernel Interface)も 0.5.0 になって、コンパイラやランタイムまで幅広く手が入ってるの。

今までどうだったの？

Neuron はこれまでもバージョンを重ねるごとに機能を増やしてきたんだけど、今回はいくつかまとまった改善が来た感じ。特にコード生成の仕組みや、細かい設定の手間が地味にボトルネックになってた部分に手が入ったのがポイントだよ。

これで何が変わるの？

まず NKI 0.5.0 では、MX FP8 のスケール型対応、テンソルの間接指定によるインデックスアクセスの命令削減、レイアウト変換がゼロコストでできる NkiTensor view API が追加されたよ。カーネルを書く人にとっては、より細かいチューニングがしやすくなったってことだね。

さらに Amazon EKS 向けに Neuron UltraServer Operator がパブリックベータで登場！Trainium の UltraServer ワークロード向けに、UltraServer の検出やワークロード割り当て、リソースクレームの生成を自動化してくれるの。EKS で大規模な Trainium クラスタを運用してる人には結構うれしい話だと思うな。

深く潜ってみよう

Neuron Compiler は、Trn2 と Trn3 でデフォルト有効になった 再設計後のコード生成バックエンド でパフォーマンスが上がったんだって。Neuron Runtime には連続共有スクラッチパッドのサポートが加わって、これまで手動で設定してたスクラッチパッドのページサイズを気にしなくてよくなったの。地味だけど運用がラクになるやつだね。

NKI Library には実験的なカーネルが 14 個 新しく追加されたよ。MoE トレーニングの集約処理、変形可能アテンション(deformable attention)、DeepSeek の MLA プロジェクション、リングアテンションなど、最近の大規模モデル開発でよく話題になる領域をカバーしてて、PyTorch のリファレンス実装も一緒に提供されてるんだって。

デバッグ面では Neuron Explorer もアップデートされて、System Trace Viewer からソースコードへのリンクや、デフォルトのグルーピングが見直されて、ワークロードのデバッグがしやすくなったよ。

対応リージョンは、Trn1・Trn2・Trn3・Inf2・Inf1 インスタンスが使えるすべての AWS リージョンだよ。

まとめ

• AWS Neuron が 2.31.0 に、NKI も 0.5.0 にアップデート
• NKI に MX FP8 対応・テンソル間接指定・ゼロコストレイアウト変換の NkiTensor view API が追加
• Amazon EKS 向け Neuron UltraServer Operator がパブリックベータで登場
• コンパイラは Trn2・Trn3 で新コード生成バックエンドがデフォルト化、ランタイムは連続共有スクラッチパッド対応
• NKI Library に実験的カーネルが 14 個追加、Neuron Explorer のデバッグ体験も改善

Trainium・Inferentia でモデルを動かしてる人、特にカーネルレベルのチューニングや EKS 上での大規模運用をしてるエンジニアに刺さるアップデートだったよ！