re:Invent 2025: NKIによるLLM最適化とNeuronでのパフォーマンスエンジニアリング実践
はじめに
海外の様々な講演を日本語記事に書き起こすことで、隠れた良質な情報をもっと身近なものに。そんなコンセプトで進める本企画で今回取り上げるプレゼンテーションはこちら！
re:Invent 2025 の書き起こし記事については、こちらの Spreadsheet に情報をまとめています。合わせてご確認ください
📖 re:Invent 2025: AWS re:Invent 2025 - Performance engineering on Neuron: How to optimize your LLM with NKI (AIM414)
この動画では、AWS TrainiumとInferentiaのAI/MLアクセラレータ向けのパフォーマンス最適化について、Neuron Kernel Interface（NKI）を用いた実践的なデモが行われています。Scott PerryとSadaf Rasoolが、Qwen3 0.6Bモデルをベースラインで実行した際の毎秒0.35プロンプト（レイテンシー約2.88秒）から、NKIで実装したflash attentionカーネルに置き換えることで毎秒2.99プロンプトへと6～8倍のパフォーマンス向上を達成する過程を示しています。NeuronCoreのTensor Engine、Vector Engine、Scalar Engineといった4つの専用コンピュートエンジンの並列活用、オンチップSRAMとHBM間のメモリ階層最適化、そしてroofline modelに基づくメモリバウンドからコンピュートバウンドへの移行戦略が解説され、昨日ローンチされたオープンソースのNKI Libraryによる事前最適化カーネルの活用方法が紹介されています。
※ こちらは既存の講演の内容を最大限維持しつつ自動生成した記事になります。誤字脱字や誤った内容が記載される可能性がありますのでご留意下さい。
本編
セッション紹介：AWS Neuron と NKI によるLLM最適化
本日のセッションへようこそ。こちらは AIM 414: Performance Engineering on Neuron—How to Optimize Your LLM with NKI というセッションです。私は Scott Perry で、AWS の Annapurna Labs チームで AI/ML パフォーマンスの Principal Solutions Architect をしています。私たちは Trainium と Inferentia という AI/ML アクセラレータを使用しているお客様と直接協力しています。本日は同僚の Sadaf が参加しています。皆さんこんにちは、Sadaf Rasool です。私も AWS Trainium で AI/ML パフォーマンスの Solutions Architect をしており、主に Neuron 上での AI/ML パフォーマンスに焦点を当てています。私たちは本日このセッションにお越しいただき、本当に嬉しいです。
ご参加ありがとうございます。これはコードトークなので、セッションの大部分をターミナルエディタとコードエディタで過ごし、スクリプトを実行し、すべてをライブで行います。うまくいくといいのですが。その前に、背景スライドをいくつか見ていく必要があります。Neuron と NKI が何であるかについて、皆さんに理解していただくためです。もう少しの間お付き合いください。その後すぐにコードに入ります。
背景をお伝えするために、長年にわたってお客様は AWS で機械学習のワークロードを実行してきました。お客様から得た初期のフィードバックの一つは、選択肢が欲しい、そして AWS で ML ワークロードを実行する際のプライス・パフォーマンスの向上を見たいということでした。2019 年頃、私たちは AWS で EC2 インスタンス経由で提供する目的別設計の機械学習アクセラレータの設計に投資しました。その最初の世代は AWS Inferentia と呼ばれていました。
これが Inferentia 1 で、EC2 インスタンスに詳しい方はご存知かもしれません。BERT や YOLO のような当時のディープラーニングモデル向けに設計された小型のアクセラレータでした。それ以来、私たちは第 2 世代の Inferentia 2 をリリースしました。さらに興味深いことに、AWS Trainium チップの 3 つのバージョンもリリースしており、第 3 世代は昨日ちょうどローンチしたばかりです。AWS Trainium は実は、分散トレーニングと推論の両方に使用できるより強力なチップです。
NeuronCore のアーキテクチャとメモリ階層の理解
Trainium と Inferentia アクセラレータチップの中核には NeuronCore があります。特定のインスタンスでは複数のチップを持つことができ、各チップ内では複数の NeuronCore を持つことができます。本日は Trn2 インスタンスで作業する予定で、これは実際にチップ内に 8 つの物理コアを持っています。ここでの例では、そのうちの 1 つを使用します。そのコア内では、これが実際に関わっているものの高レベルの図です。
このアーキテクチャで本当に興味深いのは、機械学習のワークロードのために目的別に設計されているということです。NeuronCore 内には 4 つの特定のコンピュートエンジンがあります。まず Tensor Engine があり、これは行列乗算やトランスポーズのような処理に焦点を当てています。次に Vector Engine があり、これは max pool や average pool のようなレイヤーに使用できます。Scalar Engine は活性化関数のような処理用で、さらに一般的な SIMD エンジンと汎用コアがあり、他の 3 つのエンジンではカバーされていないレイヤーのために再プログラムできます。つまり、このチップ上で実際に 4 つのコンピュートエンジンを並列に使用して、コンピュート利用率を最大化できるということです。
また注目すべき点として、オンチップの SRAM が 2 つあります：SBUF と PSUM です。これらは明らかにオンチップなので、コンピュートに非常に近く、高いバンド幅を提供しますが、容量は低めです。データを移動するための DMA エンジンのセットがあり、デバイスの一部として高バンド幅メモリもあります。これらは EC2 を通じて提供されているため、ホストメモリ、つまりこの場合は CPU に関連付けられた DRAM にもアクセスできます。つまり、メモリ階層には 3 つのレベルがあり、私たちはそれらを気にする必要があります。
オンチップ SRAM は容量は低いですが、非常に高いバンド幅のメモリで、理想的にはすべてのデータがいつもそこに存在していればいいのですが。階層の一番下には、ホストメモリがあり、容量は高いですが、バンド幅は低めです。その間に、accelerator HBM との妥協点があります：中程度の容量と、かなり良いバンド幅です。機械学習のワークロードを最適化しようとするときの大きな課題の 1 つは、データをどのように配置し、必要なときに正確にコンピュートに近づけ、階層間の移動をどのようにやりくりするかということです。
Roofline Model とパフォーマンス最適化の戦略
機械学習のワークロードを最適化することについて話すとき、与えられたワークロードを表現する 1 つの方法は roofline model です。roofline model を見たことがある人はいますか？では簡単に復習させてください。roofline model の一部として、基本的には任意のアルゴリズムまたは機械学習モデルに対して、そのモデルに内在する算術強度、つまりアルゴリズムの算術強度があります。与えられたモデルの数学的操作が何であれ、読み込まれたメモリのバイト当たりで実行する必要がある一定量の操作を消費します。しかし、与えられた accelerator チップに対して、有限のメモリバンド幅と達成できる有限の最大コンピュートスループットがあります。
与えられたワークロードに対して達成された算術強度が、グラフの左側に落ちる場合、つまり ops per byte が低い場合、メモリバウンドになります。私たちは本当にそこにいたくありません。なぜなら、それは実際に操作しているワークロードのためにメモリからどのくらい速く読み込めるかによって制限されているということを意味するからです。私たちがしたいことは、グラフの右側にものを押し出そうとして、コンピュートバウンドになることです。つまり、メモリ読み込みあたり多くの操作を実行しているので、コンピュート accelerator を完全に活用できるということです。
algorithmic な理論的な算術強度と実現された強度の違いは、基本的には実装なんです。では、パフォーマンスを改善して、グラフを右にシフトさせるにはどうすればいいでしょうか？
パイプライニング操作のようなことができます。大きなワークロードがある場合、それをより小さな部分に分割して、実際に任意の時点でワークロードの複数の部分を同時に実行し、エンジンをそれ以外の場合よりも忙しく保つことができます。データ移動の最小化のようなことがあります。多くの DMA を行っていて、不必要に大量のデータを間違ったタイミングで移動させている場合、それはパフォーマンスを低下させるでしょう。ですから、それを避けるようにしたいのです。
データスループットを最大化したいのです。場合によっては、小さなメッセージを送信したり、小さなデータを周りに移動させたりしているかもしれません。それをより大きなチャンクに統合して、データスループットを改善する方法があるでしょうか？そして、分散設定にいる場合、時々通信する必要がある複数の Trainium チップがあるかもしれません。例えば、トレーニング中に勾配を同期するような場合です。そのような場合、collectives を何らかの計算またはデータ移動と重複させることができれば、collectives が行われるのをただ待っているだけではなくなります。
Neuron SDK と NKI：低レベルカーネル開発への道
これが Trainium と Neuron にとってパフォーマンスを改善する理由と方法ですが、それは何を意味するのでしょうか？どうやってそこに到達するのでしょうか？Trainium と Inferentia チップを使用するために、Neuron SDK があります。これはソフトウェア開発キットで、スタックのさまざまなレイヤーのすべてを含んでいます。コンパイラ、ランタイム、ドライバ、そしてユーザーツールがあり、顧客がトレーニングと推論を使用できるようにしています。
通常、顧客を 3 つのペルソナとして見ています。ML 開発者、データサイエンティスト、そしてパフォーマンスエンジニアです。多くの人は基本的にその間のどこかに該当しますが、今日のところは、パフォーマンスエンジニアのペルソナにより焦点を当てています。Neuron SDK 内には、パフォーマンスエンジニアに適用される特定のツールセットがあります。今日、最も焦点を当てたいのは Neuron Kernel Interface、つまり NKI と呼ばれるものです。誰か NKI を使用したことがあるか、それについて聞いたことがある人はいますか？
では NKI とは何かというと、NKI は基本的には Trainium と Inferentia 用のカーネルを書くための Python ベースのドメイン固有言語です。NKI が登場する前は、PyTorch や JAX でマシンラーニングモデルを書いて、Neuron コンパイラを通して実行していました。Neuron コンパイラはモデルからコンピュートグラフを取り出して、それを Trainium や Inferentia で実行するように最適化して、その後モデルを実行するという流れです。ただ、汎用の Neuron コンパイラがモデルを完璧に最適化してくれるということに、かなり頼ることになってしまいます。
私たちのアクセラレータや他のアクセラレータを見ると、多くの場合、顧客はより低レベルのアクセスが必要になります。基盤となるハードウェアを最大限に活用するために、モデルを本当に細かくチューニングしたいということですね。それが NKI が提供するものです。OpenAI Triton に詳しい方であれば、これは Triton とは少し違うんですが、これは私たちのアクセラレータ用の低レベルカーネルを書くための私たちのフレーバーです。
NKI は PyTorch、JAX、NumPy と直接統合されています。Neuron ISA 命令を直接生成するので、チップに存在する実際のハードウェア命令がこのプラットフォームから生成されます。NKI 内にはいくつかの名前空間があります。nki.lang という高レベルの構成要素があって、ヘルパー関数などで皆さんをサポートします。ただし、nki.isa 名前空間の一部として直接 Neuron ISA 命令も用意されています。左側に見えるのはカーネルです。これから実際にそこに飛び込んで、行ごとに説明して、基本的なカーネルを構築するのがどのようなものかを示してから、実際のベンチマーキングに進みます。
NKI カーネルの基本構造：Tensor Add の実装例
今日のセッションでは、ここで何を達成しようとしているのかというと、記録的な短時間で達成しようとしています。まず LLM をベンチマークしてプロファイリングすることから始めます。この場合は Qwen3 になります。Hugging Face transformers と Neuron SDK を使ってそれを実行します。基本的な Qwen 実装をコンパイルして、ナイーブなアテンションを使って Trainium で実行します。これが私たちのベースラインになります。パフォーマンスと基本的なプロファイルを見てみましょう。
その後、アテンションの実際の NKI カーネル実装を説明します。それを以前実行した Qwen3 モデルに追加して、ベンチマーキングとプロファイリングをやり直して、セッションの最後にどのようなパフォーマンスの向上が得られるかを見ます。そしてもちろん、最後に皆さんからの質問の時間を設けます。今日具体的に使用しているテクノロジースタックは AWS Trainium2、つまり私たちのチップの第 2 世代です。
より小さなインスタンスでは、チップ上に存在する8つの NeuronCore のうちの1つを使用するつもりです。物事をもう少し理解しやすくするためです。Neuron SDK を使用してモデルをコンパイルして実行します。今回のケースでは、Qwen3 0.6B というエンベディングモデルを使用しており、これは Hugging Face transformers ライブラリから利用できます。
では、スライドに載せたサンプルカーネルをコードエディタで簡単に見てみましょう。それを説明してから、すぐにベンチマーキングに進んで、バニラ実装をお見せします。
ここで見ることができるのは、NKI Tensor Add kernel という非常に基本的なカーネルです。これは単に同じ形状の2つのテンソルを受け取って、それらの要素ごとの加算を行うだけの基本的なカーネルです。まず最初に気付いていただきたいのは、これが基本的な Python 関数として定義されているということです。この関数について本当に特別なことは、@nki.jit() というデコレータで装飾されているということだけで、これは Neuron コンパイラに対して、これがコンパイルしたい NKI カーネルであることを知らせます。
NKI カーネルには今日、特定の制約があります。入力と出力、つまり加算される2つのテンソルである A 入力と B 入力、そして出力ですが、これらは HBM 上に存在する必要があります。メモリの3つの層について説明しましたが、入力と出力は実際にはデバイスメモリ上に存在する必要がありますが、オンチップの SBUF 上には存在しません。フレームワークがそれを処理してくれます。今日は PyTorch を使用します。
これを1行ずつ説明していきましょう。加算したい2つの入力を渡します。同じ形状であることを確認するための基本的なチェックがあります。そうでなければ、要素ごとの加算は機能しません。ここで、すべての NKI カーネルに存在する同様のフローが見られます。最初のステップは、オンチップの SBUF 上にいくつかのスペースを割り当てる必要があるということです。ここで sbuf.view を使用して、A 入力テンソルと B 入力テンソルの両方のスペースを割り当てており、SBUF はオンチップの SRAM の1つです。
次のステップは、実際に ISA 命令を使用します。これは DMA コピーで、A 入力と B 入力の HBM テンソルをオンチップテンソルにコピーしたいことを指定します。同様に、結果用のスペースを割り当てます。これを C tile と呼び、SBUF 上にあります。sbuf.view を使ってテンソル用のスペースを割り当てて、その後 ISA の tensor_tensor 命令を使って実際に加算を行います。
ここで tensor_tensor 操作が加算演算子を使用していることを指定できます。2 つの入力は、SBUF にコピーした 2 つのタイルで、結果は出力先である C tile に出力されます。最後に、出力を HBM に戻してフレームワークに返す必要があるため、hbm.view を使って HBM ホストテンソルを割り当てて、その後再び DMA コピーを使って結果を HBM に戻してから結果を返します。
これはハイレベルなフローを示すためのかなり単純なカーネルですが、PyTorch を使い慣れている場合、これはかなり低レベルであることがわかります。ただし、本当にモデルから最高のパフォーマンスを引き出したい、ハードウェアを最大限に活用したいのであれば、このレベルまで深く掘り下げる必要がある場合があります。
Qwen3 モデルのベースライン測定とプロファイリング
では、ここでベンチマークスクリプトを簡単に見てみましょう。 スクリプトを盲目的に実行したくなかったので、実行内容を少し説明しながら進めていきます。これは Hugging Face Transformers と Qwen3 モデルを使った基本的な Python スクリプトです。 基本的に、Qwen3 モデルを CPU に最初にロードして、Neuron 用にコンパイルして、その後大量の推論を実行してレイテンシーをタイミングし、その方法でスループットを測定します。また、CPU に対してサニティチェックを行って、モデルの結果が CPU と同じであることを確認します。
インポートをスキップすると、 Neuron ランタイム用の環境セットアップが少しあります。この場合、例えば単一コアのみを使用することを指定したいです。また、プロファイリングを行うため、プロファイリング環境変数をここで有効にして、必要なメトリクスをキャプチャできるようにします。
ここでは基本的なクラスでモデルをラップしています。 これはこの埋め込みモデルから最後の隠れ状態を取得するのを少し簡単にするためだけのもので、今日のデモをちょっと楽にするためです。このクラスは本当に何もしていなくて、最後の隠れ状態だけを出力する forward パスを実行するだけです。
また、サンプルテキストをエンコードして、指定したバッチサイズに基づいて繰り返されるテンソルを返すヘルパー関数もあります。必要なバッチサイズに基づいて同じテキストを何度も繰り返すだけで、私たちを助けてくれるものです。ベンチマークスクリプトの動作を制御するために使用する引数がいくつかあります。当然のことながら、バッチサイズと最大長のようなものを制御できるのは、実行を行う際に調整できるので良いですね。 ただし、ここに NKI フラグもあります。これはベンチマークスクリプトが Sadaf が後で追加してくれる NKI 実装を使用するかどうかを指定するために使用します。
ここでは、初めてモデルを CPU にロードしている場所が見えます。 45分でこれを完了させるために、ここでモデルを4層に切り詰めています。
本質的には、少数の層のパフォーマンスは、フルモデルのパフォーマンスと似ているはずです。それに応じてスケールアップするだけです。ここから少し下にスキップすると、 ここが実際に本題に入るところです。以前にこの設定でスクリプトを実行した場合、neuron コンパイル済みモデルをディスクにキャッシュして、再度実行したい場合に時間を節約するだけです。
ただし、コンパイルされていない場合は、この torch_neuron.trace 呼び出しを使用します。これが実際に neuron コンパイラをトリガーして、コードを取得するものです。この場合、Qwen3 の vanilla 実装と naive attention だけです。それを neuron コンパイラに渡します。neuron コンパイラは、グラフを抽出して、提供したサンプル入力を使用して neuron 上で実行するようにコンパイルし、その後ディスクに保存してキャッシュできる PyTorch ロード可能なモデルを返します。
コンパイルしていなければコンパイルして、既にディスクにあれば読み込むだけです。簡単なウォームアップ推論を行って、5回の推論を実行しながら推論時間を測定して、その後 CPU との比較で精度を計算して結果を出します。Hugging Face を使ったことがあれば、これはかなり単純明快なはずですが、コードをお見せしたかったので謎のままにならないようにしました。では実際に実行して、ここから本題に入りましょう。
ターミナルを開くと、Trn2.3xlarge インスタンスにログインしていて、neuron SDK 環境がすでに利用可能になっています。例えば neuron-ls のようなコマンドを実行して、このインスタンスタイプで利用可能なものの概要をすぐに確認できます。ここでベンチマークスクリプトを Python で実行してみましょう。前の CPU 出力がディスクに保存されているのが検出されて、時間を節約するために読み込まれ、また事前コンパイルされた neuron モデルも読み込まれるはずです。数分節約できます。
数秒以内に、neuron コアでモデルを実行しているのが見え始めるはずです。これは Qwen3 0.6B で、そのモデルの 4 層バージョンで 16K のシーケンス長を使用しています。実行がほぼ完了して、精度を確認できます。MSE ロスとコサイン類似度の両方が期待通りです。アクセラレータで実行された neuron コンパイルモデルの出力は CPU 出力と 100 パーセント一致しており、これは素晴らしいことです。
パフォーマンスの面では、このモデルはそのままではあまりパフォーマンスが良くないことがわかります。1 秒あたり 0.35 プロンプトしか得られていなくて、レイテンシは約 2.88 秒です。ここには確実に改善の余地があると思います。transformer モデルのパフォーマンスを改善するために、私たちがよく注目する領域の 1 つは attention ブロックです。すぐにそれを見てみましょう。
その前に、ここでキャプチャしたプロファイリング結果も実際に確認できることをお見せしたかったのです。時々見かけるのは、実は Trainium でデバイス時間が非常に良好で、モデルが実際にうまく実行されているのに、PyTorch レイヤーで何か他のことが起こっていて、処理が遅くなっているという場合です。システムプロファイルでシステムレベルのビューを取得すれば、このモデルが最適化を追加する前に、今のところ期待できる最高の状態で実際に実行されたかどうかを確認できます。
Neuron SDK の一部として、neuron-profile というプロファイリングツールがあります。これにはさまざまなインターフェースがあります。今日は、システムレベルのプロファイルを Perfetto 互換形式で出力するだけなので、Perfetto にすぐに読み込むことができます。Perfetto はオープンソースの可視化ツールで、おそらくご存知だと思います。neuron-profile view を実行して、プロファイリングデータを含むディレクトリを指定し、Perfetto の出力形式を使用するように指定します。
ディレクトリを見ると、システムプロファイルが作成されているので、それをすぐにダウンロードします。 ブラウザに移動しましょう。これが Perfetto です。 トレースファイルを開いて、ここで簡単に見てみます。 ご存知ない場合は、これは実行タイムラインを表示しているだけで、左から始まり、右にスクロールしていきます。最初に起こったことは nrt_load です。これは実際には neuron runtime がモデルの重みを HBM に読み込んでいるのです。
それ以降は、これらすべての nrt_execute 呼び出しが表示されます。素晴らしいのは、neuron runtime 内で実際に何が起こっているかを正確に見ることができるということです。ここではモデルを実行しています。基本的にモデル内に 7 つの実行があるのが見えます。ベンチマークループは 5 でしたが、精度チェック用に 1 つ実行し、クイックウォームアップ用に 1 つ実行したので、それで 7 つになります。
これらの実行のいずれかをクリックすると、実行に関連するメタデータを実際に見ることができます。このモデルではそれほど興味深くないかもしれませんが、より大きなモデルや、複数のグラフが異なるコアで実行されている実行の場合、これは非常に役立つでしょう。 ここでは、グラフの正確な名前が見えます。これは Neuron 用にコンパイルされたモデルのバージョンで、どの NeuronCore で実行されているかが見えます。左側では、実行の期間がベンチマークスクリプトで見たものとほぼ同じであることが見えます。実行レイテンシーは 3 秒未満で、これは良くないと言えます。うまくいけば、これを改善できるでしょう。これはベースラインパフォーマンスです。Hugging Face から元の Qwen3 モデルを取得し、独自のシステムで実行するようにコンパイルし、Trainium で実行して、これらのベースライン結果を得ました。では、Sadaf が NKI を使用してこれを改善できるかどうか見てみましょう。
NKI Attention カーネルの統合：実装とコード変更
ありがとうございます、Scott。ステージを設定してくれてありがとうございます。皆さん、こんにちは。Sadaf です。昼食後にコードについて話すのは良いタイミングではないことは知っていますが、できるだけ興味深く、有望なものにしようと約束します。本質的に、ここで試みていることは、このQwen3 モデルで見たパフォーマンスを見ることです。パフォーマンス数とタイムラインも得ました。NKI、つまり Scott が言及した Neuron Kernel Interface が、このパフォーマンスを改善するのに役立つかどうか見てみましょう。
このトークのこの部分では、3つの主要な目標に焦点を当てていきます。まず、既存のモデルコードに NKI カーネルを統合することがいかに迅速で便利かをお見せしたいと思います。次に、それによるパフォーマンスへの影響は何かということです。3つ目は、拡張性とアプリケーションの観点からより考えるべき点として、皆さんが家に帰ったときに各自のワークロードで NKI カーネルを使用する方法があるかどうかということです。これらが、このトークのこの部分で取り扱う3つの主要な目標です。
少しコードを書いていくつもりです。約束した通り、退屈にならないようにしますが、ベストを尽くします。これまでのプロファイリングから注目すべき点がいくつかあります。各イテレーションについて、約2秒812ミリ秒、つまりおよそ3秒かかっています。次に、Scott が示そうとしていたように、ここで見られるスループットは毎秒0.35プロンプトです。技術的には、1つのプロンプトを処理するのに約3秒かかることを意味しています。
NKI についてそしてそれがどのようにパフォーマンスを改善するのに役立つかについて話すとき、それができる理由は、非常に低いレベルで動作し、ハードウェアに非常に近いところで動作し、基盤となるハードウェア（この場合は Trainium）から最大のパフォーマンスを得られるようにすべての利用可能な最適化を処理するからです。昨日 NKI ライブラリをローンチしたことを非常に誇りに思い、興奮しています。Scott がすぐに話すことになります。私たちは、このQwen3モデルで示すように、モデルにプラグインできるいくつかの事前定義された、事前実装されたカーネルを公開しており、Trainium 上のパフォーマンスを活用することができます。
Hugging Face Transformers ライブラリで以前に作業したことがある人なら、各モデルディレクトリに異なるモデル固有のファイルがあることをご存知でしょう。Qwen3 の場合、設定ファイル、モデリングファイルなどがあります。私が今非常に興味を持っている最初のファイルは、このモデリングファイルです。このQwen3モデリングファイルに興味を持っている理由は、Hugging Face Transformers ライブラリによって提供されるこのファイルが、Qwen3モデルの完全な実装を教えてくれるからです。Qwen3 の RMS norm がどのようなものかを理解したければ、ここにあります。Qwen3 の MLP がどのようなものかを理解したければ、ここにあります。そして rotary embedding がどのようなものか、またはそれがどのように実装されているか、すべてがここにあります。
今日、私たちが特に興味を持っている部分は attention です。誰もが transformer ベースのモデルと attention メカニズムについて聞いたことがありますよね？それが今日探索していくものです。Hugging Face Transformers ライブラリ自体によって提供されたこの eager attention forward メソッドがあり、これを NKI attention に置き換える方法があるかどうかを見ていきます。パフォーマンスへの影響を調べます。それが今日焦点を当てるモジュールです。まずこの eager attention forward をここから取りますそしてペーストします。皆さんはコードが見えていますよね？最大化する必要がありますか？大丈夫ですか？素晴らしい。eager attention をコピーしました。意味のある名前を付けましょう。NKI attention forward と呼びましょう。入力引数には触れませんが、この実装は削除します。なぜなら、ここに独自の実装を書くつもりだからです。attention 出力を返します。推論なので、重みを返す必要はありません。ですから大丈夫です。
ステップ1の目的は、既存のモデルにNKIカーネルを統合することがいかに簡単で便利かを示すことです。それが今から見ていくことです。ここにattention_forward.pyファイルがあり、その中に含まれているのはNKIの実装、異なるフレーバーのattentionの実装です。ご存知の通り、SDP attention、sliding window attention、flash attentionなど、多くのバリエーションがあります。これは大多数のそれらのattentionメカニズムの実装を持っています。その中の1つを使用しようとしています。この場合、Qwen3のためにflash attentionを活用しようとしています。
この実装ファイルが行うことは、アダプタを提供する、またはインターフェース関数と考えることができます。それは私たちからすべての入力を受け取り、正しい形状、フォーマット、データ型を与えているかどうかを確認するなどの基本的な衛生管理を確保し、その後、私たちのために正しいカーネルを呼び出します。私たちは確実にそのカーネルも見ていきますが、その前に、モデリングファイルからこのカーネルを呼び出しましょう。私のattention出力はここから来ます。そして私はこれを呼び出そうとしています。見ていたように、提供されているドキュメントを参照することができます。NKIの実装であるカーネルなので、Q、K、V テンソルを特定のフォーマットまたは特定の形状で提供することを期待しています。ハードウェアに非常に近いので、NKIにそれらのテンソルを与える前に、少しのマッサージと再形状化を行うことを期待しています。それでそれが操作を実行できるようにします。
そのために、今Q、K、V テンソルの少しの再形状化が必要です。クエリについて、最初にやることはクエリ、キー、またはバリューの形状を公開することです。この場合、バッチサイズ、ヘッド数、シーケンス長、ヘッドディメンションを取得します。それが与えてくれるものです。このカーネル、attentionカーネルは異なるフォーマットでクエリを期待しています。ここにコメントとして書き込もうとしています。それで私たちはそれを追跡し続けることができます。それは、あなたは今このフォーマットにいますが、私はあなたにこのフォーマットでクエリを与えてほしいと言っています。それはほんの少しの再形状化、permute、そしてこれらすべての古典的なテンソル操作です。ここで実行しようとしています。
ですから、私たちがやろうとしていることはB、H、S、Dからb times H、D、Sに変換することです。本質的に、2つのディメンションを1つのディメンションに縮約し、その後S、DのシーケンスをD、Sに変更しています。それだけです。最初にpermuteを実行してS、Dを切り替えるか、シーケンス長とヘッドディメンションをスワップすることができます。ですから最初にB、H、D、Sにします。そしてそこからB times H、D、Sにするのは簡単だと思います。
やってみましょう。Query equals to Query.clone() dot permute。これをやっている理由は、形状をBHSDからBHDSに変更する必要があるからです。ですから0は0のままで、1は1のままで、2は3になり、3は2になります。ここまでは良好です。私たちはBHDSを達成しました。そしてそれをB times HDSに再形状化します。非常に簡単です。そして連続したメモリロケーションを与えてくれるよう求めます。それで最適化されます。私たちはちょうどクエリテンソルを再形状化しました。今、同様に、キーテンソルに対しても同じことができます。キーは正確に同じ形状を持っているので、ここで変数名を変更するだけです。key_statesです。そして他はすべて同じままです。
値に関しては、さらにシンプルです。値に関しては、B times HSD をくれと言っています。ですから、ここで permute をする必要はありません。単純に reshape するだけで大丈夫です。permute をする必要がないので、S と D を元の位置に保つことができます。これができたら、これらのテンソルを attention kernel adapter に提供する準備ができたということです。そして、これを呼び出すので、query、key_states、value_states、そして scaling を提供できます。変数名を変更するだけで大丈夫です。key ではなく values にします。
ここまでで、私たちが行ったことは QKV の形状を変更したということです。 テンソルの値に対しては何も変更していません。他に何も変更していません。このカーネルが期待する形式に reshape しただけです。それができたら、NKI kernel adapter を通してこの attention を呼び出すことができます。パラメータを渡して、attention の出力を得ます。attention の出力を得たら、素晴らしいことに、この形式で attention の出力を与えてくれます。 そして、最終的な結果としてこの形式で持ちたいです。BSHD です。ですから、ここでも全く同じ変換をする必要があります。
こうすることができます。まず、reshape して BHSD になるようにします。そして、S と H をスワップするだけですね。今は HS を持っています。 それを SH にします。ですから、私たちのお気に入りの関数を呼び出します。それは transpose です。そして、1 番目のインデックスを 2 番目のインデックスと transpose します。Contiguous です。そして、attention の出力を返します。この実装を見ると、 query、key、values を reshape しただけです。それらを reshape したら、attention カーネルを呼び出して、出力を得ます。そして、関数の出力が要求する期待される形式で再び reshape して、 それを使用します。
もう一つやらなければならないことがあります。 attention の実装を切り替える必要があります。ここで attention を呼び出している場所があります。私がここで既に行ったことは、 これを追加しました。config に NKI が attention の実装として含まれている場合、nki_attention_forward を呼び出します。ですから、eager attention から NKI attention に切り替えます。これで、正しい attention を呼び出していることが 100% 確実になります。
素早く見てみましょう。前のプロファイルを削除します。Scott が実行していたときに見たように、彼はプロファイルを取得することができました。キャッシュデータなどがないように削除します。 そして、それを実行させてみます。ターミナルを開きましょう。 そして、今回は、Scott が実行したのと同じスクリプトを呼び出します。ただし、今回は、この引数を渡します。ですね。では、最初の実行がうまくいくことを願っています。やってみましょう。
パフォーマンス検証：6～8倍の劇的な改善を実現
これからどうなるかというと、新しい Attention の実装が見つかったので、モデルを再度 Neuron compiler に提出することになります。「新しい Attention が見つかったから、これを Trainium で実行できるようにコンパイルしてくれ」という感じですね。このコンパイルには数分かかります。その間に、約束していた実際の Attention kernel を見てみましょう。Adapter kernel を見ることで Attention forward を確認します。これらの詳細はすべて保存しておくので、後で確認できます。基本的に見えるのは、Attention kernel の実装をスワップなしで呼び出しているということです。
私がもっと知りたいのは、このカーネルの定義がどこにあるかということです。ここです、657行目にあります。ちょっと警告しておくと、これはかなりしっかりした実装で、約1000行のコードがあるので、それで退屈させたくないんです。でも共有したいのは、結局のところ Attention は本質的に QKV の関数に過ぎないということです。QK の転置、Softmax を行って、それを V と掛け算しています。これがまさにやっていることなんですが、非常にハードウェア対応的な方法でやっています。
つまり NKI を使うことで、基盤となるハードウェア、つまり Trainium の利点を最大限に活かすことができる素晴らしい機能があるということです。メモリと計算の最適化技術を採用しています。メモリ最適化技術について言うと、スマート DMA を使って DMA エンジンを効率的に使用し、DMA パケットサイズを効率的に使用しています。インテリジェントなタイリング概念、効率的なメモリレイアウト、メモリ割り当てを使用しています。計算に関しては、パイプライニングの利点を活かしているので、Scott が以前のスライドで言及していたように、異なる計算エンジンがあるので、すべての計算エンジンをできるだけ並列に忙しく保つようにして、利用可能な計算をすべて活用し、最高のパフォーマンスを得られるようにしています。
また、複数の操作を組み合わせるフュージョンの概念も使用しているので、多くの DMA や多くのメモリアクセスをする必要がなく、計算エンジンをより忙しく保つことができます。つまり、これらすべてのメカニズムを採用して、Trainium から最高のパフォーマンスを得ています。これらすべてのメカニズムを採用すると、モデルまたは特定のモジュールの超最適化された実装が得られます。
数分後、コンパイラのステータスが passed になります。これで新しい Attention 実装での推論を実行することになります。覚えておいてください、もう naive attention を実行していません。今回は NKI attention を実行しています。見えるパフォーマンスは毎秒2.99プロンプトです。naive attention で見たパフォーマンスが何だったか誰か思い出してくれませんか？そうです、0.352、ほぼ毎秒3プロンプトです。つまり、約6～8倍のパフォーマンス向上について話しています。6～8パーセントではなく、6～8倍のパフォーマンスブーストです。素晴らしくないですか？
では、プロファイルも見てみましょう。Scott が生成したプロファイルと同じように、私たちもプロファイルを生成し直します。このシステムプロファイルをダウンロードして、NKI という名前で保存することにしましょう。Profetto に戻ります。今回も覚えておいてください、Scott が 2 秒、812 ミリ秒、およそ 3 秒と言っていましたよね。では、今取得した NKI トレースを開いて、少し展開してみましょう。時間を見てみましょう。準備はいいですか？では、そこに行くと、期間がほぼ 3 秒からほぼ 4 分の 1 秒になりますよね？
この NKI kernel を使うことで、ほぼ 6 倍から 8 倍の パフォーマンスブーストが目の前で得られます。これが私の 2 番目の目標でした。1 番目は、NKI kernel がいかに簡単に統合できるかを示すことで、それは実現しました。2 番目は、パフォーマンスへの影響とそこから得られるパフォーマンスブーストを理解することでした。
3 番目の目標は、考えるための食料です。考えてみてください。これはほんの小さなモデルで、モデルの attention モジュールまたは attention 部分のうち、NKI で置き換えたのはほんの 1 つだけです。このモデルにはそれ以上にたくさんのものがあります。MLP もあり、それは超計算集約的です。QKV projections、output projections など、いろいろあります。Scott が今から私たちを導き、これらの kernel の一部が初めて外部に公開されていることを理解するのを手伝ってくれます。これらは本番規模の kernel で、私たち自身のモデルで使用でき、この素晴らしいパフォーマンスブーストを活用できます。
パフォーマンスゲインに戻ると、naive attention のスループットについては、この 1 階建てのビルからほぼ 10 階建てのビルになりました。レイテンシーの観点から見ると、できるだけ小さくしたいのですが、ここから ここまで来ました。現在のところ、消費電力に基づくプロファイリングはありません。以上で、Scott にバトンを渡します。彼がここからゴールまで導いてくれます。ありがとうございました。
NKI Library のローンチとセッションのまとめ
ありがとうございます、Sada。ご覧の通り、私たちは単にコードを数行追加し、基本的に Annapurna Labs チームによって構築されたカーネルを活用することで、あまり良くないパフォーマンスから本当に素晴らしいパフォーマンスの向上を実現することができました。ただし、私たちが本当に話したいのは、昨日 NKI Library をローンチしたということです。これはオープンソースプロジェクトで、私たちのチームによって事前に最適化された NKI カーネルが多数含まれており、私たちのチームによってキュレーション、テスト、メンテナンスされ、GitHub で提供されています。
まず最初に、今日の時点で多くのカーネルが利用可能です。QR コードをスキャンして確認することができます。最初は密集モデルのサポートが多くあります。今後、mixture of experts や他のユースケース、ワークロードの周辺にも追加のカーネルを追加していく予定です。これは本当に素晴らしいことです。なぜなら、これまで NKI はまだベータ版ですが、以前は持っていなかった低レベルのコンピュートエンジンへのアクセスを顧客に提供するからです。注意メカニズムの素朴な実装をコンパイルしたときの様子を見ました。実行され、正確な結果が得られましたが、すぐに出た状態ではパフォーマンスは本当に冴えませんでした。
NKI を使用することで、そこのギャップを大幅に縮めることができ、このモデルで素晴らしいパフォーマンスを実現することができました。今日は注意メカニズムに焦点を当てましたが、異なるモデルと異なるレイヤーの他の多くの側面があり、NKI を適用できる可能性があります。これは少しオニオンピーリングのようなものです。ワークロードに入り、デバイスレベルのプロファイルを実行し、ボトルネックがどこにあるかを把握し、NKI でそれらのいくつかに取り組むことになります。ボトルネックが他の場所に移動し、ある時点で、ほぼ永遠にそれを追い続けることができます。しかし、少なくとも NKI があれば、顧客の手にこれらの機能が与えられるので、PyTorch レベルからは以前は不可能だったこのパフォーマンスエンジニアリングを実際に行うことができます。
それでは、本日ご来場いただいた皆様にお礼申し上げます。午後のセッションで Trainium と NKI について学ぶために来ていただき、本当にありがとうございます。本日と明日は Neuron と Trainium に関する他の多くのセッションがあり、興味があればチェックしてみてください。ご質問があればお受けします。ありがとうございました。
※ こちらの記事は Amazon Bedrock を利用し、元動画の情報をできる限り維持しつつ自動で作成しています。
Discussion