TileFuse: A Fused Mixed-Precision Kernel Library for Efficient Quantized LLM Inference on AMD NPUs

本論文は、AMD XDNA2 NPU向けの、アンパッキング、デクオンタイズ、および行列演算を融合させることで、普及しているAWQスタイルの量子化LLM推論に対する効率的なネイティブサポートを可能にし、クライアントエッジデバイス上の既存のベースラインに対して大幅なパフォーマンス向上とエネルギー効率の向上を実現する、ハードウェアに近い（close-to-metal）混合精度カーネルライブラリであるTileFuseを導入するものである。

要約

ガレージに、新品で超高速な配送トラック（NPU）が置いてあると想像してください。しかし、手元にある唯一の地図（ソフトウェア）は、遅い古い自転車の乗り方しか教えてくれません。トラックのスピードを活用することはできません。なぜなら、その地図はトラックの走行レーンをどのようにナビゲートすべきかを知らないからです。

これが、今日のノートPC用チップ上で最新のAIチャットボット（LLM）を動かす際に直面している問題です。これらのチップには、バッテリーを節約し高速に動作するように設計された強力な「AIエンジン（NPU）」が搭載されていますが、通常、ソフトウェアはエンジンに合わせてAIの形を変えるよう強制してしまいます。エンジンがAIに適応させるのではなく、AIがエンジンに合わせる形です。

TileFuseは、これを解決する新しいツールセットです。これは、トラック専用のカスタム高速道路を建設するようなものです。これにより、データを再パッケージ化するために立ち止まることなく、圧縮された大量のデータを運ぶことができるようになります。

以下に、この論文がこれらの解決策をシンプルな比喩を用いてどのように説明しているかを記します。

1. 問題点：「再パッケージ化」のボトルネック

通常、これらの新しいチップでAIを実行するには、AIの「圧縮された」重み（これは、小さなスーツケースにきつく詰め込まれた本のようなものです）を、チップが読み取れるように嵩高い形式（本を取り出してテーブルの上に平らに広げるようなもの）へと展開（アンパック）しなければなりません。

• 従来の方法： チップはスーツケースを読み取り、本を取り出し、別のスーツケースに入れ直してから、ようやく読み取りを開始します。これは時間とエネルギーの無駄です。
• TileFuseの方法： チップはスーツケースを読み取り、中身を展開しながら、そのまま本を読み取ります。すべてを一つのスムーズな動きで行います。

2. 解決策：「融合（Fused）」カーネル

著者らは TileFuse というライブラリを作成しました。「カーネル」とは、チップに対する特定の指示書のようなものだと考えてください。

• 融合（Fusion）： 3人の別々の作業員（展開する人、変換する人、計算する人）を用意する代わりに、TileFuseは彼らを一人のスーパーワーカーに統合します。このワーカーは、圧縮されたデータを掴み、その場で変換し、計算までを一度に行います。
• 結果： これは、単に野菜を切るだけでなく、切る、味付けする、調理するという工程を、一つの連続した動作で行うシェフのようなものです。論文によれば、これにより特定のタスクにおいて、従来のメソッドと比較して「展開」の部分が最大 2.8倍速く なります。

3. 「インターリーブ（Interleaved）」レイアウト：倉庫の整理

大規模なAIモデルは、膨大な数字のリスト（重み）を持っています。チップのメモリシステムには、次のデータを取りに行くために到達できる距離に限界があります。もしデータがバラバラで乱雑に保存されていると、チップは次のピースを手に入れるために長く、遅い移動を繰り返さなければなりません。

• 比喩： 倉庫の中で箱がランダムに積み上げられている様子を想像してください。フォークリフトは次の箱を取るために50フィートも走らなければなりません。
• TileFuseによる修正： 彼らは、次に必要な箱がすぐ隣にあるように、倉庫の配置（「インターリーブ・プリタイリング」と呼ばれます）を再編成しました。これにより、チップは大量のデータを一度の滑らかなスウィープでまとめて掴むことができ、以前は適合できなかった非常に大きなAIモデル（最大32,000アイテム幅）のサポートが可能になりました。

4. 「GEMV」問題：交通渋滞

AIチャットボットは2つのフェーズで動作します。

1. プリフィリング（Prefilling）： 長いプロンプトを一気に読み込むこと（本を一冊丸ごと読むようなもの）。これはトラックにとって速く、簡単な作業です。
2. トークン生成（Token Generation）： 一度に一つの単語を書いていくこと（文章を書いていくようなもの）。これは遅く、トリッキーです。

• 課題： 一度に一つの単語を書いているとき、チップの「トラック」はしばれて動いています。なぜなら、それは大きな荷物を運ぶように設計されており、小さな荷物には向いていないからです。これは、一通の手紙を届けるためにセミトレーラーを使用しているようなものです。エンジンはかかっていますが、トラックの中身は空っぽなのです。
• 解決策： TileFuseはこのフェーズの交通の流れを再設計しました。データを高速道路のたった一つのレーンに送るのではなく、チップの32レーンすべてに同時に仕事を分散させます。これにより、たとえ「荷物」が小さくても、エンジン全体を稼働させ続けることができます。

5. 実世界の成果

チームは実際のAMDノートPC（Ryzen AI）でテストを行い、標準的なグラフィックスカード（iGPU）と比較しました。

• 速度： 長いプロンプトの読み込み（プリフィリング）において、TileFuseを備えたNPUは、グラフィックスカードよりも最大 2倍高速 でした。
• バッテリー： NPUはより効率的であるため、同じ作業を行うのに 64%少ないエネルギー しか消費しませんでした。
• 注意点： 一度に一つの単語を書くとき（トークン生成）、NPUは時としてグラフィックスカードよりも遅くなりました。これは、NPUを構成するための「セットアップ時間」が、このような極めて小さく素早いタスクに対しては長すぎるためです。
• ハイブリッド・ソリューション： 論文は「両方の良いとこ取り」のアプローチを提案しています。つまり、重い作業（長いプロンプトの読み込み）にはNPUを使い、素早いタスク（単語の書き出し）にはグラフィックスカードを使うという方法です。この組み合わせが、最高のスピードとバッテリー寿命をもたらします。

まとめ

TileFuse は架け橋です。開発者がすでに使用している一般的な圧縮AIフォーマット（AWQなど）を取り込み、AIモデル自体を変更することなく、AMDの新しいAIチップ上でネイティブに動作するようにします。展開ステップと計算ステップを融合させ、データを完璧に整理し、チップの全レーンを活用することで、ノートPCでのAI実行を大幅に高速化し、エネルギー効率を高めます。

技術概要

技術要約: TileFuse

問題提起

エッジデバイスへの大規模言語モデル（LLM）のデプロイは、厳格な電力および熱設計枠を満たすために、ニューラル・プロセッシング・ユニット（NPU）への依存を強めています。しかし、クライアント向けNPU、具体的にはAMDのXDNA2アーキテクチャにおける実用的なLLM推論は、大きな障壁に直面しています。主な障害は、ソフトウェアスタックの断片化と低レベルなプログラマビリティの欠如です。ほとんどのNPUフレームワーク（例：Apple Core ML、Intel OpenVINO、Qualcomm QNN）は、高レベルなインターフェースや独自のデプロイパスのみを公開しており、AWQ（Activation-Weight Quantization）のような広く使用されている既成の量子化フォーマットをネイティブにサポートしていません。

その結果、開発者はモデルをNPU固有の量子化スキームに適合するように再構成するか、あるいはより効率の低いフル精度での実行に頼らざことは避けられません。先行研究の中には、特定のプラットフォームに対してハードウェア認識型の量子化を探索したもの（例：QualcommのScalingLLM）もありますが、それらは多くの場合、$W4A16$のような標準的なフォーマットをサポートしていないか、現代のLLMに共通する大きな行列次元（例：8K超）を扱うことができません。さらに、既存のNPU最適化は、アンパッキング（展開）とデクオンタイゼーション（逆量子化）を計算カーネルに直接融合させる能力に欠けていることが多く、不要なメモリトラフィックや非効率なデータ移動を引き起こしています。

手法

TileFuseは、AMD XDNA2 NPUを対象とし、LLM推論を支配する線形層（GEMMおよびGEMV）をターゲットとした、ハードウェアに近い（close-to-metal）混合精度カーネルライブラリです。本システムは、オープンなMLIR-AIE（IRON）プログラミングフローを活用し、ウェイトのレイアウト、メタデータの配置、マイクロカーネル、およびアレイレベルのデータフローを協調設計（co-design）します。

1. 融合された混合精度マイクロカーネル

TileFuseは、3つの異なる操作を単一のカーネルフローに融合させます：

• アンパッキング（Unpacking）: パックされたINT4/INT8ウェイトの抽出。
• デクオンタイゼーション（Dequantization）: パーグループまたはパーチャネルのスケールおよびゼロポイントを用いた、低ビット整数からBF16への変換。
• 実行（Execution）: 行列乗算（GEMM）または行列ベクトル乗算（GEMV）の実行。

ウェイトを計算コアに到達するまで圧縮形式で保持することで、TileFuseはフル精度（BF16）のウェイトを実体化するために必要なグローバルメモリへのラウンドトリップを回避します。マイクロカーネルは、各ウェイトタイルを一度だけアンパックおよびデクオンタイズし、その結果得られたBF16値をローカルの計算コアメモリに保存し、後続の行列演算のために再利用するように設計されています。これにより、タイル型GEMMの性能に不可欠な再利用パターンが維持されます。

2. インターリーブ型プレタイリング（Interleaved Pre-Tiling）

ランタイムのDMAストリーミング制約（通常、サポートされる次元を8K未満に制限する）の限界に対処するため、TileFuseはオフラインのインターリーブ型プレタイリング戦略を採用しています。

• メタデータパッキング: 量子化メタデータ（スケールおよびゼロポイント）をメモリ内のウェイトタイルに直接パッキングし、デクオンタイゼーションに必要なすべてのデータが連続して利用可能であることを保証します。
• レイアウト再編成: ウェイト行列をインターリーブされた列主体のレイアウトに再編成します。これにより、同じAIEカラムに割り当てられたウェイトタイルがメモリ上で連続することになり、単一のGEMM操作が複数のカーネル呼び出しに分割される原因となる大きなメモリストライドを排除します。このアプローチにより、最大32KまでのGEMM次元をサポートします。

3. GEMVデータフローの協調設計

エッジLLMにおけるトークン生成は、通常、バッチサイズ1のGEMV操作ですが、これはフル精度ではメモリバウンドになりますが、量子化とデクオンタイゼーションが導入されると計算バウンドになります。

• ベースラインの限界: 標準的なGEMVマッピングは、通常、4×8のAIE計算アレイの単一の行のみを利用するため、ハードウェアの大部分がアイドル状態になります。
• 最適化: TileFuseは、ウェイトをフル4×8アレイ全体に分散させるようにデータフローを再設計します。メモリコアを中間配布層として利用し、シムコアからメモリコアへウェイトタイルの束をストリーミングし、そこから4つの計算行へとファンアウト（拡散）させます。これにより、空間的利用率と算術スループットを最大化し、トークン生成を実現します。

4. 再構成可能性を考慮した最適化

本システムは、XDNA2アーキテクチャの再構成可能性を明示的に活用し、特定のワークロードに合わせて計算コアのマイクロカーネルとDMAアクセスパターンをカスタマイズします。ただし、著者らはこの柔軟性がランタイムのディスパッチおよびファブリック再構成のオーバーヘッドを導入することも認めています。TileFuseは、大規模なプリフィリングGEMMにおいてこのオーバーヘッドを償却するように最適化されていますが、レイテンシに敏感な小規模なGEMV操作における限界も認識しています。

主な貢献

1. 実用的な量子化のための融合カーネル: AWQスタイルの$W4A16$および$W8A16$ウェイトを直接消費する融合AIEカーネルの実装により、モデルの再形成や独自の量子化スキームの必要性を排除しました。
2. インターリーブ型プレタイリング: GEMM次元を最大32Kまでサポートし、メモリアクセスの整合性をDMAバッファ記述子の制約に合わせることで、効率的なランタイムストリーミングを可能にするレイアウト戦略。
3. GEMVデータフローの再設計: ウェイトをフル4×8 AIEアレイに分散させる新しいデータフローにより、トークン生成のスループットを大幅に向上。
4. エンドツーエンドのエッジ効率: これらのカーネルをハイブリッドランタイム（線形層にはNPU、アテンション/非線形層にはiGPUを使用）に統合し、レイテンシとエネルギー効率の大幅な向上を実証。

結果

評価は、Llama3-8B、Gemma-2B、Qwen2.5-3Bなどのモデルを用い、Ryzen AI ラップトップ（Ryzen AI 7 350 および Ryzen AI 9 HX 370）上で実施されました。

• カーネル性能:
  • GEMM: TileFuseは、フル精度ベースラインに対して最大**121.6%**の性能向上を達成。
  • GEMV: TileFuseは、フル精度ベースラインに対して最大**281%**の性能向上を達成。
  • 比較: GEMMにおいて、TileFuseは強力なiGPUベースライン（llama.cppの融合カーネル）に対して、2倍以上の性能およびエネルギー効率の向上を提供。
• エンドツーエンドのLLM推論:
  • プリフィリング（Prefilling）: Ryzen AI 7 350において、TileFuseはiGPUベースラインと比較して、最大2.0倍低いプリフィリングレイテンシ（例：Llama3-8B、2048トークンにおいて2.00倍の高速化）を達成。
  • エネルギー: Ryzen AI 7 350においてエネルギー消費を64.6%以上削減し、Ryzen AI 9 HX 370では**52.8%**削減。
  • トークン生成: トークン生成においては、マイクロ秒単位の実行時間に対するミリ秒単位のカーネルディスパッチおよび再構成のオーバーヘッドのため、NPUはiGPUを下回ります。
• 先行研究との比較: TileFuseは、Llama3-8BおよびQwen2.5-3Bにおいて、それぞれ1.45倍および1.61倍の高速化を達成し、FastFlowLMやScaling NPUといった最近のNPU駆動型フレームワークを上回るプリフィリングスループットを示しました。

重要性と主張

本論文は、TileFuseが、AMD XDNA2 NPUが効率的で既成の量子化LLM推論にとって実用的なターゲットであることを示していると主張しています。AWQのような標準的なフォーマットを、低レベルな（close-to-metal）最適化を通じて直接ハードウェアに持ち込むことで、本研究は、生のNPUの潜在能力と現実世界のクライアント展開との間の溝を埋めています。

著者らは、TileFuseの成功はハイブリッドデプロイ戦略に依存していることを強調しています。すなわち、大規模なコンテキストのプリフィリング（オーバーヘッドが償却される場面）にはNPUの再構成可能な空間アレイを利用し、レイテンシに敏感なトークン生成にはiGPUに頼るという戦略です。このアプローチは、モデルを独自の制約に従わせるのではなく、ハードウェアをユーザーが依存する量子化フォーマットに適応させれば、NPUがエッジAIの性能とエネルギーの面で大幅な利益をもたらし得ることを検証しています。本研究は、適切なカーネルおよびデータフローの協調設計があれば、ソフトウェアスタックの断片化が存在する場合でも、NPUがエッジAIにおいて顕著な性能とエネルギーの恩益を提供できることを確立しています。