Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator

本論文は、LLM サービングにおける動的な KV キャッシュ管理と分散メモリ抽象化のギャップを埋め、ハイブリッドボンディング技術を活用した「Helios」というハードウェア・ソフトウェア協調設計のアクセラレータを提案し、既存の GPU や NMP 設計と比較して大幅な速度向上とエネルギー効率の改善を実現したことを述べています。

要約

🚀 背景：AI は「メモリの壁」にぶつかった

まず、現在の AI 生成サービス（チャットボットなど）が抱えている問題から始めましょう。

• 問題点: AI は「計算が得意な頭脳」と「大量のデータを覚える記憶」の両方を必要とします。
• 現状のボトルネック: 従来の GPU（グラフィックボード）は「計算」は超高速ですが、「記憶（メモリ）」からのデータ読み取りが少し遅いのです。
  • 例え話: 天才的な料理人（計算エンジン）がいても、冷蔵庫（メモリ）が遠く離れていて、食材を取りに行くのに時間がかかると、料理を作るスピードは遅くなります。
• AI の特徴: AI は会話が進むにつれて、過去の会話内容（文脈）をすべて覚えておく必要があります。この「過去の会話データ」はKV キャッシュと呼ばれます。ユーザーの質問の長さや、同時に何人かの人が使っているかによって、このデータの量や動き方が刻一刻と変わります（非常に動的です）。

これまでの「近接メモリ処理（NMP）」という技術は、冷蔵庫を料理人の隣に近づけようとしましたが、**「固定された棚」**しかありませんでした。

• 古い方式の欠点: 「1 つの棚には、1 人のユーザーの全データを丸ごと入れる」と決まっていたため、ユーザーが少ないときは棚が空っぽで無駄になり、ユーザーが増えたり会話長くなったりすると、棚が足りなくなってパニックになります。

---

💡 Helios の解決策：「ハイブリッド・ボンディング」と「柔軟な棚」

Helios は、この問題を 2 つの大きな工夫で解決します。

1. 技術の進化：「ハイブリッド・ボンディング」

これは、「メモリ（冷蔵庫）」と「計算チップ（料理人）」を、3 次元に重ねて、まるで「溶接」のように密着させる技術です。

• メリット: 距離が極端に短くなるので、食材（データ）の移動が光速のように速くなり、エネルギーもほとんど消費しません。

2. 発想の転換：「動的なブロック管理」

これが Helios の最大の特徴です。

• 古い方式（固定棚）: 「1 人のユーザーには、最大限のスペースを 1 つの棚に割り当てる」。
• Helios の方式（ブロック管理）: 「データを小さな**『ブロック（箱）』**に分けて、空いている棚に自由に散らばって配置する」。

🍕 ピザの例え:

• 古い方式: 100 人分のピザを、100 枚の大きなトレイに 1 枚ずつ並べる。でも、実際に注文が 10 人しか来ないと、90 枚のトレイが空っぽで場所を占領している。
• Helios: ピザを「1 切れずつ」の小さな箱に分ける。注文が来たら、空いている箱を 10 個集めて渡す。注文が増えたら、また空いている箱を足す。
  • 結果: どのユーザーが来ても、どの長さの会話でも、「空いているスペース」を最大限に有効活用できます。

---

🤝 協調作業：「分散タイル・アテンション」

データをバラバラに配置すると、計算するときに「あっちの棚からこっちの棚へデータを集める」手間がかかります。Helios はこれを賢く処理します。

• 仕組み: 複数の小さな計算ユニット（PE）が、まるで**「チームワークでパズルを解く」**ように動きます。
• 工夫:
  • データを「タイル（タイル状のブロック）」に分けて、各ユニットが並行して計算します。
  • 計算中に「オンライン・ソフトマックス」という処理を同時に行い、データのやり取りの待ち時間をゼロにします。
  • 例え話: 大人数で大きな壁紙を貼る作業を想像してください。古い方式だと「1 人が壁の全部を貼る」のを待ってから次の人が始めますが、Helios は「1 人ずつ小さなタイルを貼っていき、最後にパズルのようにつなぐ」ことで、全員が同時に働き、待ち時間なしで完了させます。

---

📊 結果：どれくらい速くなった？

Helios を従来の GPU や他の NMP 技術と比較した結果、驚異的な性能向上が確認されました。

• 速度: 平均して3.25 倍速くなりました。
• 省エネ: 平均して3.36 倍のエネルギー効率になりました。
• 待ち時間: ユーザーが「次の言葉が出るまで」待つ時間が、最大で**72%〜76%**も短縮されました。

特に、**「急に人が増えた時」や「長い会話をする時」**に、他の技術が詰まってしまうのに対し、Helios はスムーズに動き続けることができます。

---

🌟 まとめ

Helios は、「AI の記憶（メモリ）」と「計算（プロセッサ）」を密着させ、さらにデータを「小さな箱」に分割して柔軟に配置することで、AI チャットボットを「待ち時間なし・低コスト・高効率」で動かすための画期的な設計図です。

まるで、**「固定された棚」から「空いているスペースを賢く見つけて使う、流動的な倉庫」**へと進化させたようなイメージです。これにより、未来の AI サービスは、もっと速く、もっと安価に、そしてもっと多くのユーザーに提供できるようになるでしょう。

技術概要

論文「Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator」の技術的サマリー

本論文は、大規模言語モデル（LLM）のオンライン推論サービスにおける高い動的なワークロード（可変なリクエスト長、変動する到着率）に対応するための、3D-DRAM（ハイブリッドボンディング技術）を活用した新しいアクセラレータ「HeLioS」を提案しています。既存の Near-Memory Processing (NMP) 設計の限界を克服し、ハードウェアとソフトウェアの共設計（Co-design）によって、LLM サービングの性能とエネルギー効率を大幅に向上させることを目指しています。

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

---

1. 問題定義 (Problem)

LLM のオンライン推論サービスでは、以下のような高い動的な特性が求められますが、既存の NMP 型アクセラレータには以下の根本的な課題が存在します。

• LLM サービングの動的性質:
  • リクエスト長の多様性: 文脈長（コンテキスト）が平均 1K トークンから最大 10K トークン以上まで大きく変動する。
  • 到着率の変動: ピーク時とオフピーク時で 1〜2 桁の差があり、バッチサイズも小規模から大規模まで変化する。
  • KV キャッシュの管理: 推論プロセス（プレフィルとデコーディング）において、KV キャッシュはリクエストのライフサイクルに依存して動的に増減し、ブロック単位で管理される必要がある。
• 既存 NMP 設計の限界:
  • 粗粒度な KV キャッシュ管理: 既存の設計では、アテンションヘッドの KV キャッシュ全体を固定された NMP モジュール（例：HBM クューブ）に静的に割り当てる。これにより、リクエスト長が短い場合のメモリ断片化や、負荷の偏りが発生し、実効的なサービング容量が低下する。
  • 柔軟性の欠如したアテンション実行フロー: 既存の NMP は、アテンション計算を 2 つの独立した GEMM（行列積）として実行し、KV キャッシュを固定されたモジュール内で完結させる必要がある。これにより、ブロック単位の KV キャッシュ管理や、動的なリクエスト長への対応が困難になっている。
  • アーキテクチャとアルゴリズムのミスマッチ: NMP の非均一なメモリ抽象化（各処理エンジン PE がローカルメモリしか直接アクセスできない）と、GPU ベースの動的 KV キャッシュ管理アルゴリズム（全メモリへの自由なアクセスを前提）との間に根本的なギャップがある。

2. 手法と提案システム (Methodology & HeLioS)

著者は、HeLioS（Hybrid-bonding-based LLM Serving accelerator）を提案し、以下の 3 層の設計で上記の問題を解決します。

A. アーキテクチャ設計 (Architecture)

• ハイブリッドボンディング (Hybrid Bonding) の活用:
  • DRAM デルとロジックチップを Cu-Cu 直接結合で 3D 積層。高い I/O 密度（110,000/mm²）と低抵抗により、HBM や GDDR よりも優れた帯域幅とエネルギー効率（0.66 pJ/bit）を実現。
  • 4 層の DRAM 積層により、GPU に匹敵する大容量（≥80GB）を確保しつつ、ロジックチップ上でカスタム化された計算エンジン（PE）を配置。
• 分散タイルド・アテンション (Distributed Tiled Attention):
  • 各 PE がローカルメモリにアクセスし、アテンション計算を「タイル（ブロック）」単位で分散実行。
  • PE 間通信はメッシュトポロジーの NoC を経由し、部分和の集約やデータ転送を最適化。

B. 演算実行フロー (Operator Execution)

• PE 内実行:
  • オンラインソフトマックス（Online Softmax）と部分和集約を、GEMM 計算とオーバーラップさせることでレイテンシを隠蔽。
  • プレフィル段階の KV キャッシュ転送には、転送バッファと計算バッファを分離し、メモリアクセス遅延と計算遅延に応じて「細粒度」または「粗粒度」の転送戦略を動的に選択。
• PE 間通信:
  • アテンションブロック: クエリベクトルの全 PE への複製（All-Gather）、キー/バリューベクトルの宛先 PE への分散（Scatter）を、メッシュトポロジーに最適化された集合通信プリミティブ（Reduce-Scatter, All-Reduce など）で効率的に実行。
  • MLA (Multi-Head Latent Attention) 対応: 潜在 KV キャッシュを共有し、行列吸収（Matrix Absorption）技術を用いて中間結果の膨大なメモリ使用を回避。

C. システム設計 (System Design)

• 空間認識型 KV キャッシュ割当 (Spatially-Aware KV Cache Allocation):
  • アルゴリズム 2: リクエストの KV ブロックを、PE 間の計算負荷バランスと通信オーバーヘッドの最小化を考慮して割り当てる。
  • 割当基準:
    1. 各 PE の総トークン数（計算/ストレージ負荷）を均等化。
    2. 通信コストを最小化するため、中央の PE（Reduce/Scatter のハブとなる）に近い位置には「最後のブロック」を割り当て、遠い位置には「中間ブロック」を割り当てる戦略を採用。
  • これにより、動的なワークロード下でも PE 間の負荷偏りを抑制し、メモリ断片化を最小限に抑える。

3. 主要な貢献 (Key Contributions)

1. 既存 NMP 設計の限界の深掘り分析: 動的な LLM サービングワークロードにおいて、既存の NMP が粗粒度管理と固定的な実行フローにより非効率であることを実証。
2. HeLioS アーキテクチャの提案: ハイブリッドボンディング技術を用いた、NMP ネイティブな分散タイルド・アテンションを可能にするハードウェア設計。
3. NMP ネイティブな動的 KV キャッシュ管理: 非均一なメモリ抽象化を考慮した、空間認識型のブロック割当アルゴリズムと、分散実行フローの共設計。
4. 広範な評価: 多様なモデル（OPT, LLaMA3, Mixtral, Qwen3, DeepSeek）とワークロード条件下での評価。

4. 評価結果 (Results)

既存の GPU および NMP ベースの設計（AttAcc, NeuPIMs, Duplex, Stratum など）と比較して、HeLioS は以下の成果を達成しました。

• 性能向上:
  • 平均（幾何平均）で 3.25 倍 の速度向上。
  • 既存の GPU/NMP 設計と比較して、最大で 72% (P50) / 76% (P99) のトークン間時間（TBT）の短縮。
• エネルギー効率:
  • 平均で 3.36 倍 のエネルギー効率の向上。
• メモリ断片化の低減:
  • 既存 NMP 設計に比べ、ピーク時のメモリ断片化比率が 2.21 倍〜10.60 倍 低減。これにより、高負荷時でも高いサービング容量を維持可能。
• 負荷分散:
  • PE 間の計算負荷の偏り（最大 - 最小の遅延ギャップ）を 7.1%〜13.3% 以内に抑え、既存設計（19.9%〜66.6%）よりも大幅に改善。
• WaferLLM との比較:
  • 同様のメッシュ型アクセラレータである WaferLLM と比較し、長文脈ワークロードにおいて 1.55 倍 の計算速度向上と 6.11 倍 の通信効率向上を実現。

5. 意義 (Significance)

本論文の提案する HeLioS は、LLM サービングの「動的性」と NMP アーキテクチャの「分散性」の間の根本的なミスマッチを解決する重要なステップです。

• 実用的なサービング能力の向上: 変動するリクエスト長や到着率に対応できる柔軟な KV キャッシュ管理により、クラウド環境での LLM 推論コストを大幅に削減し、スループットを向上させます。
• 次世代ハードウェアの指針: ハイブリッドボンディング技術を活用した 3D-DRAM アーキテクチャが、メモリバンドルット制約の厳しい LLM 推論において、従来の GPU 以上のポテンシャルを持つことを示しました。
• システム共設計の重要性: ハードウェア（PE 配置、通信トポロジー）とソフトウェア（ブロック割当アルゴリズム、実行フロー）を密接に連携させることで、単なるハードウェアの高速化を超えた性能向上が可能であることを実証しました。

結論として、HeLioS は、大規模な LLM 推論サービスにおいて、高効率かつ低遅延を実現するための有望な次世代アクセラレータ設計として、ハードウェアとソフトウェアの両面から重要な貢献を果たしています。