手動カーネルチューニングはもう終わり：Neuron Agentic DevelopmentがAWS Trainiumの最適化を加速する方法

フロンティアAIモデルの規模と複雑さが増すにつれ、開発者はすべてのハードウェアプラットフォームで共通の課題に直面しています。それは、モデルが実行されるシリコンから最大のパフォーマンスと効率を引き出す方法です。ワールドモデルのリアルタイム体験、エージェントワークフローでの深い推論のサポート、大規模推論コストの削減など、ハードウェアが理論的に提供できるものと、ほとんどのチームが実際に達成できるものとの間には依然として大きなギャップがあります。カスタムカーネル開発は伝統的にそのギャップを埋める手段でしたが、深いアーキテクチャ知識、手動プロファイリングワークフロー、反復的な最適化サイクルが必要であり、ほとんどのチームには負担が大きすぎます。

本日、AWSはNeuron Agentic Development機能を発表します。これは、AWS TrainiumおよびAWS Inferentia上でビルドする開発者向けに、AIエージェントとスキルのコレクションであり、これを可能にします。最初の機能により、KiroおよびClaudeのコーディングエージェントがNeuron Kernel Interface（NKI）カーネルの作成、デバッグ、プロファイリングを行えるようになり、MLパフォーマンスエンジニアリングをチームのすべての開発者に拡大します。他のアーキテクチャから来たカーネル開発者はTrainiumに迅速に適応でき、チームはアイデアからハードウェア最適化実装までの時間を短縮でき、かつてカーネル開発をゲートしていた深いアーキテクチャ知識が、各ステップで開発者をガイドするエージェンティックツールを通じてアクセス可能になります。

本記事では、Neuron Agentic Development機能がカーネル開発ワークフローをどのように加速するかを説明します。

Neuron Agentic Developmentスキル

Neuron Agentic Developmentパッケージは、自然なカーネル開発パイプライン（作成→デバッグ→プロファイリング→分析）に従う5つの専門スキルを提供します。個別のタスクにはスキルを単独で呼び出したり、neuron-nki-agentでチェーンして、リクエストに基づいて適切なワークフローを自動選択できます。使用するには、エージェンティックIDEのスキルディレクトリにスキルを追加します。スキルはTrainiumベースのAmazon EC2インスタンスで実行する必要があります。

カーネル作成：neuron-nki-writingスキルはNKIカーネル作成の出発点です。PyTorch、NumPy、または自然言語の説明を正しいNKIコードに変換します。例えば、ハードウェア制約を尊重するタイリング戦略、メモリアクセスパターン、明示的なdstパラメータによる計算操作、DMAサイジングとSBUF再利用の効率ガイドラインをカバーします。このスキルはタスクを複雑度で分類し、必要な参照のみをロードします。

デバッグ：neuron-nki-debuggingスキルは、TrainiumおよびInferentiaハードウェア上でのNKIコンパイルおよび実行エラーを解決するための体系的なワークフローを提供します。例えば、正しい--targetフラグを使用した環境セットアップ、すべての28のNCCエラーコードのカテゴリ別インデックスによるコンパイラエラー解決、CPU計算参照に対する数値検証をカバーします。

プロファイリングと分析：neuron-nki-profilingスキルはハードウェア上で実行プロファイルをキャプチャします。ランタイム検査環境変数を設定し、カーネルを実行し、正しいNEFFを識別し、DMAレベルの詳細を含むneuron-explorerでトレースをキャプチャします。JSONメトリクスを抽出し、neuron-nki-profile-queryingが消費するNEFFファイルを生成します。neuron-nki-profile-queryingスキルはNEFFおよびNTFFファイルを取り込み、SQLクエリを実行してパフォーマンスバウンドを計算し、ボトルネックエンジンを特定し、非効率を特定のNKIソース行に特定します。3つの分析アプローチをサポート：neuron-explorer APIサーバー、Parquet上のDuckDB直接、またはカスタム計算用のpandas。

ドキュメント：neuron-nki-docsスキルは開発全体を通じて使用されます。作成中はAPIシグネチャとチュートリアルを提供し、デバッグ中はエラーコードを説明し、プロファイリング中はハードウェアアーキテクチャの詳細を明確にします。

エージェント

スキルが個々のタスクのビルディングブロックを提供する一方、エージェントは複数のスキルを自律的なワークフローに組み合わせます。各エージェントは、多段階の開発シナリオをエンドツーエンドで処理する専門化されたペルソナです。neuron-nki-agentはNKI開発の統合エントリポイントで、リクエストに基づいて自動的に適切なワークフロー（作成、デバッグ、プロファイリング、ドキュメントルックアップ）を選択し、適切なスキルを調整します。これがデフォルトの開始点です。その他に、カーネル作成に特化した作成エージェント、コンパイラエラーを自律的に解決するデバッグエージェント、軽量ドキュメントナビゲータのドキュメントエージェント、2つのスキルを組み合わせてパフォーマンスボトルネックを特定するプロファイル分析エージェントがあります。

実践：カスタムsoftmaxカーネルの最適化

以下のウォークスルーでは、これらのエージェンティック機能が実際にどのように連携するかを示します。softmaxカーネル（ステップ1、2）とSwiGLUカーネル（ステップ3、4）の2つを探索し、現実のワークロードでのプロファイリングを実演します。

推論パイプラインでPyTorch softmax操作がボトルネックになっており、先行するスケール操作と融合するカスタムNKIカーネルを作成したいとします。

ステップ0：インスタンスと環境のセットアップ——trn2.3xlargeインスタンスを起動し、KiroとNeuron Agentic Developmentスキルをインストールし、Neuronデバイスが可視であることを確認します。

ステップ1：カーネルを作成——Kiro CLIセッションで、次のプロンプトを入力します：「softmax(x * scale)を最後の次元に沿って計算するNKIカーネルを作成し、入力形状は[batch, seq_len, hidden_dim]、bfloat16。」エージェントは、ハードウェアアクセラレートされたexpを使用した3パスカーネル（行最大値、指数和、正規化）を生成し、float32累積で数値安定性を確保し、適切なタイリングを行います。

ステップ2：ハードウェア上でデバッグ——エージェントにカーネルを実行し、PyTorch参照と数値的一致を検証するよう依頼します。エージェントは即座に問題に遭遇します：nisa.tensor_tensorはリダクション結果を自動ブロードキャストしないため、行ごとの最大値と和の値を行列全体に直接適用できません。エージェントはNKI参照パターンを参照し、正しいブロードキャストメカニズム（ステップ0アクセスビュー）を特定し、カーネルを書き換えます。修正後、4つのテストケースすべてがbfloat16許容範囲内の最大誤差で合格します。

ステップ3：カーネル実行をプロファイルおよび分析——このステップでは、SwiGLU MLPカーネルを使用してプロファイリングを実演します。エージェントはカーネルをNEFFにコンパイルし、neuron-explorerでNTFFトレースをキャプチャした後、2部構成の調査を実行します。まず、キャプチャしたプロファイルで完全なバウンド分析を実行し、複数の改善余地を発見します。テンソルエンジンが実行を支配し、非効率的です。また、DMAエンジンに冗長で非効率的な作業があります。調査の結果、NKI matmul命令はすでにピーク効率に近いのに対し、DMA命令はターゲットサイズを大幅に下回っており、すべての入力を8回再ロードしていることがわかりました。非効率転送の原因となっているNKIコードの正確な3行も特定されました。これらに対処することで、テンソルエンジンのアイドルギャップを減らし、カーネルレイテンシを改善できる可能性があります。

注意事項

プロファイリングおよびデバッグスキルは実際のTrainiumまたはInferentiaベースのインスタンスでの実行が必要です。作成およびドキュメントスキルはどこでも動作します。すべてのスキルは現在のNKI API（0.4.0）を対象としています。

Neuron Agentic Developmentはカーネル開発を民主化し、より多くの開発者がチップの専門家でなくても特定のハードウェア向けにモデルを最適化できるようにしています。