PIM-SHERPA: Software Method for On-device LLM Inference by Resolving PIM Memory Attribute and Layout Inconsistencies

本論文は、オンデバイス LLM 推論におけるメモリ属性とレイアウトの不一致という課題を解決し、DRAM ダブルバッファリングやオンライン重み再配置といったソフトウェア手法により、PIM 対応システムでの効率的な推論を実現する「PIM-SHERPA」を提案するものである。

要約

この論文は、**「スマホや小型の端末で、巨大な AI（大規模言語モデル）を高速に動かすための新しい『ソフトウェアの魔法』」**について書かれたものです。

タイトルにある「PIM-SHERPA」は、この問題を解決するソフトウェアの名前です。Sherpa（シェルパ）とは、ヒマラヤの山でガイド役をする人たちのことで、AI が重い荷物を背負って登るのを助ける役割を果たします。

以下に、専門用語を避け、日常の例え話を使って簡単に解説します。

---

1. 背景：AI が抱える「二つの顔」という問題

スマホで AI を動かそうとすると、AI は大きく分けて 2 つの作業を行います。

1. 最初の入力（プリフィル）： ユーザーが長い文章を入力したとき。これは「計算が大量に必要」な作業です。
2. 次の言葉の生成（デコード）： 1 文字ずつ答えを返していくとき。これは「メモリの読み書きが大量に必要」な作業です。

ここで、**「PIM（メモリ内処理）」**という新しい技術が登場します。これは、データが保存されている「メモリ」の中に小さな計算機を内蔵し、データを持ち運ばずにそこで計算させる技術です。これにより、特に「次の言葉の生成」が劇的に速くなります。

しかし、ここに大きな問題が 2 つあります。

• 問題 A：「冷蔵庫」か「倉庫」か？（メモリの属性の矛盾）

  • 最初の入力では、同じデータを何度も使うので、**「冷蔵庫（キャッシュ）」**に入れておくと速いです。
  • 次の言葉の生成では、PIM を使うために、あえて**「冷蔵庫に入れない（キャッシュ不可）」**状態にする必要があります。なぜなら、冷蔵庫に入っていると、PIM が「あ、データがあるから読み取る必要ないね」と判断してしまい、PIM の計算が起動しないからです。
  • 矛盾： 1 つの AI モデルが、同じデータに対して「冷蔵庫に入れてほしい」と「入れないでほしい」という相反する要求を同時にしています。

• 問題 B：「本棚の並べ方」の違い（レイアウトの矛盾）

  • スマホの CPUは、本を「縦に並べる」のが得意です。
  • PIMは、本を「横に並べる」方が得意です。
  • 矛盾： どちらの作業も同じデータを使うのに、並べ方が違うと、作業をするたびに本を並べ替える（レイアウト変換）必要があります。

2. 昔の解決策：「コピーを 2 つ作る」

これまでの解決策は、**「2 つのコピーを作る」**というものでした。

• 冷蔵庫用（CPU 用）の並べ方と、PIM 用（PIM 用）の並べ方のデータを、それぞれ 1 セットずつ用意するのです。
• デメリット： メモリ（RAM）の容量が2 倍必要になります。スマホはメモリが限られているので、これだと「重い AI は動かせない」か、「もっと小さい（性能の低い）AI に我慢する」しかありませんでした。

3. 新しい解決策：PIM-SHERPA（シェルパ）

この論文では、**「コピーを 2 つ作らずに、必要な時だけ必要な場所に運ぶ」**というソフトウェアの工夫を提案しています。

方法 1：DDB（ダブルバッファリング）＝「2 つの作業台を交互に使う」

• 仕組み： 小さな「作業台（バッファ）」を 2 つ用意します。
  • 今、作業台 A で計算をしている間、裏で作業台 B に「次の計算に必要な本（データ）」を並べ替えて運んでいます。
  • 計算が終わったら、作業台 A と B を入れ替えます。
• メリット： 運ぶ時間（並べ替え）と計算する時間が重なり合うので、待ち時間がほとんど発生しません。

方法 2：OWR（オンライン並べ替え）＝「必要な時だけ急いで並べ替える」

• 仕組み： 入力する文章が長い場合、計算に時間がかかるので、その間に「必要な本だけ」をその場で並べ替えてから計算します。
• メリット： 複雑な同期制御が不要で、実装が簡単です。長い文章を扱う現代の AI には特に有効です。

4. この技術のすごいところ

1. メモリ容量を約半分にする：

   • 従来の「コピー 2 個」方式に比べて、必要なメモリ容量を約 48% 削減できました。
   • これにより、スマホでもより高性能な AI モデル（例：Llama 3.2 の 3B モデルなど）を、PIM の恩恵を受けながら動かせるようになります。

2. 速度はそのまま、性能は最大限：

   • メモリを節約しつつ、理論上の最高速度に近いパフォーマンスを維持しています。

3. ハードウェアいらず：

   • 新しいチップを作る必要はありません。既存のスマホや端末のソフトウェアをアップデートするだけで実現できます。

まとめ

この論文は、**「AI がスマホで動くとき、メモリの『冷蔵庫』と『倉庫』の使い分けや、本棚の並べ方の矛盾を、賢い『運搬係（シェルパ）』がソフトウェアだけで解決し、メモリ容量を節約しながら超高速化を実現した」**という画期的な成果を報告しています。

これにより、今後私たちが使うスマホの AI は、より賢く、より速く、そしてより多くの機能を搭載できるようになるでしょう。

技術概要

PIM-SHERPA: 論文の技術的サマリー

この論文は、エッジデバイス（スマートフォンなど）における大規模言語モデル（LLM）の推論を効率化するためのソフトウェア手法「PIM-SHERPA」を提案するものです。メモリ内処理（PIM: Processing-in-Memory）技術を活用する際、プリフィル（Prefill）フェーズとデコード（Decode）フェーズの間で生じるメモリ属性の不整合と重みレイアウトの不整合という根本的な課題を、ハードウェア変更なしに解決することを目的としています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

---

1. 問題定義：PIM 導入における 2 つの矛盾

LLM 推論は、計算集約的な「プリフィルフェーズ」と、メモリ帯域幅集約的な「デコードフェーズ」で構成されます。PIM（特に LPDDR-PIM）はデコードフェーズのボトルネック解消に有効ですが、両フェーズを単一のシステムで効率的に実行する際に以下の 2 つの矛盾が発生します。

1. メモリ属性の不整合 (Memory Attribute Inconsistency)

   • プリフィル: 重みの再利用頻度が高いため、キャッシュ可能領域 (Cacheable) に配置してキャッシュヒット率を最大化したい。
   • デコード: PIM 実行を確実にトリガーするためには、DRAM コントローラーへのリクエストが必要であり、キャッシュ不可能領域 (Non-cacheable) に配置しなければならない（キャッシュヒットすると PIM が発動しないため）。
   • 課題: 同一の重みがフェーズによって矛盾するメモリ属性を要求する。

2. 重みレイアウトの不整合 (Weight Layout Inconsistency)

   • ホスト側（CPU/GPU）: 列方向または行方向の連続したアドレス（ホストフレンドリーなレイアウト）を好む。
   • PIM 側: 各 DRAM バンク内でデータが連続しており、SIMD 演算を最大化するための「PIM 感知型レイアウト（PIM-aware layout）」を必要とする。
   • 課題: 最適なメモリ配置がフェーズによって異なる。

既存手法の限界:

• 重みの複製 (Weight Duplication): 両方のレイアウト/属性を持つコピーを保持するが、DRAM 容量をほぼ倍増させ、エッジデバイスの容量制約（例：12GB）を突破し、モデルサイズを制限する。
• ハードウェア依存: 一部の手法（FACIL など）はメモリコントローラーの改造を必要とし、製品レベルのシステムでは適用困難。

---

2. 手法：PIM-SHERPA

PIM-SHERPA は、ハードウェア変更を伴わず、ソフトウェアのみで上記の矛盾を解決する 2 つのアプローチを提案します。これらは、インタラクティブな LLM や RAG（検索拡張生成）の普及により入力シーケンス長が増加しているトレンドに対応しています。

2.1 核心技術：Swizzled Memory Copy (SMC)

PIM 感知型レイアウト（非キャッシュ領域）から、ホストフレンドリーなレイアウト（キャッシュ領域）へデータを転送・変換するプロセスです。

• 重みの要素座標に基づき、DRAM アドレスを計算し、非キャッシュ領域からキャッシュ領域へデータをコピーしながらレイアウト変換（スウィズリング）を行います。
• これにより、標準的な GEMM カーネルをそのまま使用でき、ポータビリティを維持します。

2.2 アプローチ 1: DRAM Double Buffering (DDB)

• 仕組み: 小さなキャッシュ可能バッファ（通常は 2 つ）を確保し、現在のレイヤーの計算（GEMM）と並行して、次のレイヤーの重みを非キャッシュ領域からバッファへプリフェッチします。
• 動作:
  • バッファ 0 で計算中、バッファ 1 へ次のレイヤーの重みをコピー（SMC）する。
  • 次のレイヤーへ移行する際、バッファを切り替える。
• 利点: 計算時間とデータ転送時間を重畳（オーバーラップ）させることで、レイアウト変換のオーバーヘッドを隠蔽します。
• 適用: 入力シーケンス長が中程度の場合に有効。

2.3 アプローチ 2: Online Weight Rearrangement with SMC (OWR)

• 仕組み: 計算直前に、必要な重みのみをオンデマンドで SMC により変換・コピーします。
• 動作: 計算スレッドとコピースレッドを直列に実行します（複雑な同期制御が不要）。
• 利点: 実装がシンプルで、入力シーケンス長が十分に長い場合（GEMM 計算時間が SMC 転送時間を支配する場合）、オーバーヘッドが相対的に小さくなります。
• 適用: 長い入力シーケンス（RAG や長いコンテキスト）に適しています。

---

3. 主要な貢献

1. 問題の特定: PIM 対応 LLM 推論における「プリフィルとデコード間のメモリ属性および重みレイアウトの不整合」を初めて体系的に特定し、製品レベルのシステムにおける課題として提示しました。
2. 容量オーバーヘッドの定量化: 従来の「重み複製」アプローチが、代表的な LLM（Llama 3.2 など）において DRAM 容量をほぼ倍増させ、スマートフォンでの実行を不可能にすることを示しました。
3. ソフトウェアのみによる解決策の提案: 既存のハードウェアを変更することなく、小さなキャッシュバッファと SMC を用いて両方の不整合を解決する DDB と OWR を提案しました。

---

4. 実験結果

Samsung Galaxy S24+ (Exynos 2400) と Llama 3.2 (1B, 3B) モデルを用いた評価を行いました。

• DRAM 容量の節約:
  • 重み複製（WD）と比較して、PIM-SHERPA は約 47.8% 〜 49.7% の DRAM 容量を節約しました。
  • これにより、Llama 3.2 3B モデルを 12GB DRAM を搭載するスマートフォンで PIM 機能付きで実行可能になりました（複製では 6GB 追加が必要で不可能）。
• 推論性能 (TTFT - Time to First Token):
  • DDB: 入力シーケンス長が 128 以上の場合、理論的な最大値に近い性能を達成し、ハードウェア改造が必要な FACIL-O と同等の性能を示しました。
  • OWR: 入力シーケンス長が長い場合（例：192 トークン）、DDB や FACIL-O と同等の性能（約 16.7 TPS）を達成しました。
  • PIM による高速化: PIM-SHERPA は、PIM 未使用のベースラインと比較して最大3.3 倍の高速化を実現しました。
• 制約環境での有効性:
  • DRAM 容量が限られ、重みの複製ができない環境（3B モデルを 1 回分しか保持できない場合）でも、PIM-SHERPA は機能し、PIM 未使用の非キャッシュ GEMM（NC-GEMM）に比べて劇的な性能向上を示しました。

---

5. 意義と結論

• 実用性の向上: PIM-SHERPA は、ハードウェアの変更を必要としないソフトウェアソリューションであるため、既存のモバイルおよびエッジデバイスプラットフォームへの即時導入が可能です。
• PIM 普及の鍵: 本論文は、PIM 技術を実際の製品（特に容量制約の厳しいモバイルデバイス）で LLM 推論に適用する際の最大の障壁であった「属性とレイアウトの矛盾」を解決し、PIM の商業化に向けた重要なステップを提供します。
• 将来展望: 入力シーケンス長の増加傾向（RAG や長いコンテキスト処理）と相性が良く、将来的に計算能力の高いシステム（GPU/NPU）においても、入力シーケンス長の増加に伴い本手法の有効性は維持・向上すると予測されます。

要約すると、PIM-SHERPA は、メモリ制約の厳しいエッジデバイスにおいて、PIM の恩恵を受けながら大規模言語モデルを効率的に実行するための、画期的なソフトウェア基盤技術です。