【ITニュース解説】HeMA-MISO: Heterogeneous Memory Architecture for LLM Inference with SW Optimization

大規模言語モデル（LLM）の推論プロセスは、コンピューターのメモリシステムに大きな負担をかける。特に、大量の計算を必要とするため、多くの場合はデータがメモリから読み込まれるのを待つ「メモリ律速」の状態になりやすい。また、LLMは非常に大規模なデータを扱うため、通常のキャッシュメモリ（SRAM）が効率的に使われにくいという課題がある。さらに、Mixture of Experts（MoE）と呼ばれるモデルや、非常に長い文章（コンテキスト）を扱うLLMでは、単にメモリ容量が不足してしまう問題も深刻である。この研究では、特にメモリ容量の制約に焦点を当て、その解決策として、異なる種類のメモリを組み合わせたハードウェア構造と、それを最大限に活用するためのソフトウェア最適化を提案する。

近年のMoEモデルは、その大部分が「エキスパート」と呼ばれる専門的なサブモデルで構成されており、これが膨大なメモリ容量を必要とする。また、LLMが対話や長文の理解など、より複雑なタスクをこなすために、扱うコンテキスト長が飛躍的に伸びていることもメモリ容量の負担を増大させている。このような容量不足を解決するために、単純にAIチップ（NPU）の数を増やす方法は、計算能力の無駄やデータ転送のボトルネックを引き起こす。また、DDRやLPDDRといった大容量メモリを直接NPUに追加すると、その低いデータ転送速度が推論の効率を大きく下げてしまう。そこで、高速だが容量の小さいHBM（High Bandwidth Memory）と、大容量だがHBMよりは低速なDDR/LPDDRメモリを組み合わせた、階層的なメモリシステムが有望な解決策として考えられている。

この研究で提案されるハードウェア構造は「HeMA（Heterogeneous Memory Architecture）」と称され、HBMとDDR/LPDDRという二層のメモリ構成を特徴とする。既存のNPUがHBMを使用しているところに、さらにDDRまたはLPDDRメモリを追加する。そして、これら異なる種類のメモリを効率的に管理するため、「メモリハンドラーコア（MHC）」という専用の制御回路を導入する。計算ユニットがメモリへのアクセスを要求すると、MHCはそのアドレスを基に、高速なHBM向けか、大容量のDDR/LPDDR向けかを判断し、適切なメモリにリクエストを振り分ける。また、HeMAアーキテクチャでは、キャッシュラインと呼ばれるデータ単位で、これら二つのメモリ階層間でデータを転送できる。ほとんどのNPUにおいて、計算ユニットは「スクラッチパッドメモリ（SPM）」と呼ばれる内部の高速な一時記憶領域に存在するデータしか直接アクセスできないという前提も、この設計の重要な点である。

HeMAアーキテクチャを最大限に活用するために、「MISO（Memory-Informed Software Optimization）」というソフトウェア最適化手法も提案されている。この手法は二つの重要な観察に基づいている。一つは、LLMの推論処理ではデータがあまり繰り返し使われないため、NPUの高速なSPMが効率よく利用されていないこと。もう一つは、システムが低速なDDR/LPDDRメモリからデータを読み込んでいる間、高速なHBMやSPMが「遊んで」いる状態であることだ。MISOは、この非効率性を解消することを目的とする。

従来の推論処理では、DDRに保存されたデータの一部をSPMに移動させて計算し、その間に次のデータをDDRから読み込み始めるという流れになる。しかし、DDRのデータ転送速度が遅いため、SPMでの計算が終わっても次のデータがまだ到着しておらず、計算ユニットがデータの到着を待つ「ストール」状態が発生し、貴重な計算サイクルが無駄になる。この待ち時間中に高速なHBMは何もせずアイドル状態になっている。

MISOは、このHBMのアイドル時間を活用する。具体的には、SPMを二つの部分に分割する。SPMのごく一部をDDRから読み込んだ現在の計算に使い、SPMの残りの大部分を、次に必要となるデータを高速なHBMから「プリフェッチ（先読み）」するために利用する。このプリフェッチは、DDRからの計算と並行して行われるため、データの転送にかかる時間を隠蔽できる。DDRからの計算が完了すると、DDR用に予約されていたSPM領域も解放され、HBMからプリフェッチされたデータのためにSPM全体が利用可能になる。

このプリフェッチ手法を効果的に機能させるには、三つの条件を満たす必要がある。まず、DDRベースの計算に利用できるSPM領域を制限しても、その計算自体の性能がほとんど低下しないこと。次に、モデルのデータをメモリ上に配置する際、DDRからの計算とHBMからの計算が交互に行われるように配置し、両方のメモリシステムが効率的に利用される機会を作ること。そして最後に、DDR用のSPM領域を可能な限り小さくして、HBMからのプリフェッチに使えるSPM領域を最大化すること。これを実現するために、計算処理の最適化技術（オペレーターフュージョンやタイリング）をDDRベースの計算に積極的に適用する。

MISOの有効性を検証するために、仮想的なHeMA-TPUアーキテクチャ（TPUv4に128GBのDDRメモリを追加）を想定し、NPUシミュレータ上で実験を行った。Mixtral 8x7BモデルをFP8精度で利用し、長いコンテキスト長と大きなバッチサイズでの性能を評価した。データの配置戦略としては、KVキャッシュをHBMに優先的に置く「KV-major」、エキスパートウェイトをHBMに優先的に置く「Expert-major」、そしてMISO独自の戦略の三つを比較した。MISO戦略では、各層においてKVキャッシュとエキスパートウェイトの一部を両方ともHBMに配置することで、Attention操作とMoE操作の両方でHBMとDDRを同時に利用できるようにする。これにより、MISOが目指すデータプリフェッチが可能になる。

Attention操作の評価では、SPM容量を8MBまで制限しても性能劣化がほとんどないことを確認した。これに基づき、Attention計算の最終段階でSPMの使用を制限し、HBMからのエキスパートウェイトのプリフェッチを開始する戦略が考案された。FFN操作については、オペレーターフュージョンを適用し、SPM制限による性能低下はAttentionほどではないものの発生するが、許容範囲内であると結論付けた。最終的なデータフローは、DDRからのAttention計算中にHBMからMoEエキスパートウェイトをSPMにプリフェッチし、MoE計算中にHBMから次の層のKVキャッシュをSPMにプリフェッチするというサイクルを形成する。このSPM制限とプリフェッチは、DDR上のデータからHBM上のデータに切り替わる直前の部分にのみ適用される。

性能評価では、MISOが他の手法（最適化されたKV-majorおよびExpert-major）と比較して、一貫して最も高いスループットを達成した。ただし、MISOによる性能向上は、Flash Attentionのような基本的な最適化手法が適用されていないベースラインと比べると控えめであり、MISOの利点は既存の最適化技術の上に積み重なる形で現れる。また、入力のシーケンス長やバッチサイズが大きくなるにつれて、MISOによる性能向上の幅は縮小する傾向が見られた。これは、SPMの容量が固定されているため、総KVキャッシュサイズが増大するにつれて、プリフェッチの効果が相対的に小さくなるためである。非常に大規模なLlama4-Scoutモデルでの評価でも同様の傾向が確認され、極端に大きなKVキャッシュサイズではMISOによる性能向上は数パーセントに留まった。

レイテンシの内訳を分析すると、DDRベースのAttentionやMoE操作においては、MISOはSPMサイズを制限しているため、Expert-majorよりもわずかに遅くなることが判明した。しかし、HBMベースのAttentionやMoE操作では、MISOがデータを事前にSPMにプリフェッチしているため、Expert-majorよりも低いレイテンシを達成する。このHBM側でのレイテンシ削減が、最終的な全体スループットの向上に貢献している。

この研究にはいくつかの限界点もある。一つは、前述のように、入力サイズが大きくなるにつれてMISOの性能優位性が低下することである。もう一つは、非常に小さなバッチサイズの場合、MISOのデータ配置戦略（すべての層でHBMにエキスパートのスペースを確保する）が非効率になる可能性がある点だ。もしバッチサイズが小さすぎて、確保されたエキスパートの一部が全く利用されない場合、高性能なHBMのスペースが無駄になり、全体の性能が低下する可能性が考えられる。