KVPR: Efficient LLM Inference with I/O-Aware KV Cache Partial Recomputation

要約

大規模な言語モデル（LLMS）の推論は計算的に要求があります。
自動回帰デコードのコストを削減するために、キー値（kV）キャッシュを使用して中間活性化を保存し、トークン生成の計算オーバーヘッドを大幅に低下させます。
ただし、KVキャッシュに必要なメモリは急速に成長し、GPUメモリの容量を超えることがよくあります。
費用対効果の高い代替手段は、KVキャッシュをCPUメモリにオフロードすることです。これにより、GPUメモリ圧力が軽減されますが、ボトルネックはCPUとGPUの間のPCIE接続の限られた帯域幅にシフトします。
既存の方法は、I/OとGPU計算を重複するか、CPU-GPUの不均一な実行を採用することにより、これらの問題に対処しようとしますが、それらは過剰なデータの動きとCPU機能への依存によって妨げられます。
KVキャッシュのサイズが大きくなったり、GPU計算機能が増加するにつれて、PCIE通信のレイテンシが完全に重複します。
このホワイトペーパーでは、CPUが最初にアクティブ化の部分的なセットを最初に転送し、そこからGPUがKVキャッシュ値の再計算を開始できるKVPRを紹介します。
GPUは部分KVキャッシュを再構成しますが、KVキャッシュの残りの部分はCPUから同時に転送されます。
このアプローチは、KVキャッシュ転送とGPUの再構成を重複させて、アイドルGPU時間を最小限に抑え、推論パフォーマンスを最大化します。
KVPRは、入力特性とシステムハードウェア情報、計算および通信ワークロードの分布を最適化するスケジューラモジュール、および派生した実行計画を効率的に実行するランタイムモジュールを利用するプロファイラーモジュールを統合することにより、完全に自動化されます。
実験結果は、KVPRが最先端のアプローチと比較して、デコード中に最大35.8％低いレイテンシと46.2％高いスループットを達成することを示しています。
このコードはhttps://github.com/chaoyij/kvprで入手できます。

要約(オリジナル)

Inference for Large Language Models (LLMs) is computationally demanding. To reduce the cost of auto-regressive decoding, Key-Value (KV) cache is used to store intermediate activations, which significantly lowers the computational overhead for token generation. However, the memory required for the KV cache grows rapidly, often exceeding the capacity of GPU memory. A cost-effective alternative is to offload KV cache to CPU memory, which alleviates GPU memory pressure, but shifts the bottleneck to the limited bandwidth of the PCIe connection between the CPU and GPU. Existing methods attempt to address these issues by overlapping GPU computation with I/O or employing CPU-GPU heterogeneous execution, but they are hindered by excessive data movement and dependence on CPU capabilities. Fully overlapping PCIe communication latency gets challenging as the size of the KV cache grows and/or the GPU compute capabilities increase. In this paper, we introduce KVPR, an efficient I/O-aware LLM inference method where the CPU first transfers a partial set of activations, from which the GPU can start recomputing the KV cache values. While the GPU recomputes the partial KV cache, the remaining portion of the KV cache is transferred concurrently from the CPU. This approach overlaps GPU recomputation with KV cache transfer to minimize idle GPU time and maximize inference performance. KVPR is fully automated by integrating a profiler module that utilizes input characteristics and system hardware information, a scheduler module to optimize the distribution of computation and communication workloads, and a runtime module to efficiently execute the derived execution plan. Experimental results show that KVPR achieves up to 35.8% lower latency and 46.2% higher throughput during decoding compared to state-of-the-art approaches. The code is available at https://github.com/chaoyij/KVPR.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google