LLM Inference in a Flash!
Abstractの概要
本論文では、大容量ストレージとニアデータ実行を実現する一方で、書き込み耐性の低さ、オフチップ帯域幅の狭さ、高精度浮動小数点サポートの弱さといった制約を抱えるフラッシュベースのコンピュートインメモリ（CIM）アーキテクチャに、LLMデコーディングを適合させる手法を検討しています。これらのハードウェアボトルネックに対処するため、著者らは浮動小数点のオフロードを行わずに線形および非線形演算子を実行するエンドツーエンドの整数専用量子化パイプラインと、辞書ベースのKVキャッシュ圧縮戦略を提案しています。この圧縮手法は、過去のキーと値をフラッシュ常駐の静的辞書に対するスパースコードとして表現しつつ、直近のローカルウィンドウを高密度形式で保持し、クエリ認識型の階層的スパース性を適用します。Llama-3.1-8BおよびQwen-2.5-7Bでの評価において、本手法はタスク精度を維持し、解析的システムモデルにより、NPUベースラインと比較してCIM-SSDアーキテクチャ上でデコード遅延とエネルギーの大幅な削減が見込まれることが示されています。
新規性
主な新規性は、GPUのメモリ階層ではなく、フラッシュコンピュートインメモリの制約に特化して調整されたアルゴリズムとハードウェアの協調設計にあります。本手法は、すべての演算子にわたる整数専用実行パイプラインと、計算負荷の高い最小二乗再フィッティングを置き換える射影ベースのスパースコーディング探索を組み合わせることで、書き込みトラフィックと動的なKVキャッシュデータ移動を最小限に抑えています。
成果
Llama-3.1-8Bにおいて、整数専用パイプラインはFP16ベースラインのWikiText-2パープレキシティ7.5454に対し7.5803を達成し、クリティカルなデコードパスから浮動小数点演算を排除しました。統合された圧縮手法は、Llama-3.1-8BとQwen-2.5-7Bの双方においてFP16フルKVのLongBenchスコアの99%を維持しながら、動的KVキャッシュトラフィックを15倍削減します。解析的システムモデリングによると、デフォルトのS2 CIM-SSD構成は、DRAMベースのNPUベースラインと比較して、1Kコンテキスト長で3.1倍のレイテンシ削減と2.7倍のエネルギー削減を達成し、256Kコンテキスト長では4.4倍のレイテンシ削減と6.8倍のエネルギー削減にスケールすることが予測されています。
論文の注目点
- エンドツーエンドの整数専用LLM量子化パイプラインは、アクティベーション平滑化、グループワイズ量子化、INT16再量子化を活用して安定性を維持しつつ、Softmax、RMSNorm、SiLUの浮動小数点演算を整数近似に置き換える。
- フラッシュ常駐辞書を用いたKVキャッシュ圧縮手法は、射影ベースの探索とクエリ認識型の階層的リファインメントを用いて過去のコンテキストをスパースコードとして表現し、高コストな最小二乗解法を回避してフラッシュの書き込み摩耗を最小化する。
- WikiText-2およびLongBenchベンチマーク全体で精度の低下はごくわずかであり、解析的ハードウェアモデリングによりCIM-SSDプラットフォーム上の長文コンテキスト推論において最大4.4倍のレイテンシ削減と6.8倍のエネルギー削減が予測される。