XQuant: KVキャッシュの再計算によるLLM推論のメモリ壁の打破

要旨

LLM推論は多くの下流アプリケーションにとって重要なワークロードとして浮上しているが、LLMを効率的に推論することは、膨大なメモリフットプリントと帯域幅要件のため困難である。並行して、計算能力は過去数十年にわたりメモリ容量と帯域幅を着実に上回っており、この傾向は現代のGPUハードウェアでも顕著であり、LLM推論の課題をさらに悪化させている。そのため、メモリ操作を削減するために計算量を増やす新しいアルゴリズムが登場している。この流れを活用し、我々はXQuantを提案する。XQuantは、最先端のKVキャッシュ量子化手法と比較して大幅な精度向上を伴う低ビット量子化により、メモリ消費量を桁違いに削減する。これを実現するために、標準的なKVキャッシュの代わりに層入力アクティベーションXを量子化してキャッシュし、推論中にKeysとValuesをその場で再生成する。これにより、KVキャッシュと比較して即座に2倍のメモリ節約が得られる。XQuantを適用することで、FP16ベースラインと比較して<0.1のパープレキシティ劣化で最大約7.7倍のメモリ節約を達成する。さらに、我々のアプローチは、X値が層間で類似しているという事実を活用する。この観察に基づいて、X埋め込みの層間類似性を利用して極端な圧縮を実現するXQuant-CLを導入する。様々なモデルにおいて、XQuant-CLはFP16ベースラインと比較してわずか0.01のパープレキシティ劣化で最大10倍のメモリ節約を達成し、0.1のパープレキシティ劣化で12.5倍のメモリ節約を実現する。XQuantは、ハードウェアプラットフォームの急速に増大する計算能力を活用してメモリボトルネックを解消し、最先端のKVキャッシュ量子化手法を凌駕し、幅広いモデルでFP16に近い精度を達成する。

English

Although LLM inference has emerged as a critical workload for many downstream applications, efficiently inferring LLMs is challenging due to the substantial memory footprint and bandwidth requirements. In parallel, compute capabilities have steadily outpaced both memory capacity and bandwidth over the last few decades, a trend that remains evident in modern GPU hardware and exacerbates the challenge of LLM inference. As such, new algorithms are emerging that trade increased computation for reduced memory operations. To that end, we present XQuant, which takes advantage of this trend, enabling an order-of-magnitude reduction in memory consumption through low-bit quantization with substantial accuracy benefits relative to state-of-the-art KV cache quantization methods. We accomplish this by quantizing and caching the layer input activations X, instead of using standard KV caching, and then rematerializing the Keys and Values on-the-fly during inference. This results in an immediate 2times memory savings compared to KV caching. By applying XQuant, we achieve up to sim 7.7times memory savings with <0.1 perplexity degradation compared to the FP16 baseline. Furthermore, our approach leverages the fact that X values are similar across layers. Building on this observation, we introduce XQuant-CL, which exploits the cross-layer similarity in the X embeddings for extreme compression. Across different models, XQuant-CL attains up to 10times memory savings relative to the FP16 baseline with only 0.01 perplexity degradation, and 12.5times memory savings with only 0.1 perplexity degradation. XQuant exploits the rapidly increasing compute capabilities of hardware platforms to eliminate the memory bottleneck, while surpassing state-of-the-art KV cache quantization methods and achieving near-FP16 accuracy across a wide range of models.

XQuant: KVキャッシュの再計算によるLLM推論のメモリ壁の打破

XQuant: Breaking the Memory Wall for LLM Inference with KV Cache Rematerialization

要旨

Support