SwiftKV: 知識保存モデル変換を用いた高速プリフィル最適化推論

SwiftKV: Fast Prefill-Optimized Inference with Knowledge-Preserving Model Transformation

October 4, 2024

著者: Aurick Qiao, Zhewei Yao, Samyam Rajbhandari, Yuxiong He

cs.AI

要旨

一般的な企業での要約、RAG、およびコード生成などの人気のあるユースケースに対するLLM推論は、通常、生成長よりも桁違いに長いプロンプト長を観察します。この特性は、プリフィルの高コストと応答遅延の増加につながります。本論文では、プロンプトトークンの処理時間とコストを削減しつつ生成トークンの高品質を維持するために特に設計された新しいモデル変換および蒸留手法であるSwiftKVを提案します。SwiftKVは、次の3つの主要メカニズムを組み合わせています：i）SingleInputKVは、後のレイヤーのKVキャッシュを、はるかに早いレイヤーの出力を使用してプリフィルし、プロンプトトークンがモデル計算の大部分をスキップできるようにします。ii）AcrossKVは、隣接するレイヤーのKVキャッシュをマージしてメモリフットプリントを減らし、より高いスループットのためのより大きなバッチサイズをサポートします。iii）既存のLLMをSwiftKV用に適応させるための知識保存蒸留手法は、最小限の精度影響と低い計算およびデータ要件で実現できます。Llama-3.1-8Bおよび70Bにおいて、SwiftKVはプリフィルの計算要件を50％削減し、KVキャッシュのメモリ要件を62.5％削減し、幅広いタスクにおいて最小の品質劣化をもたらします。最適化されたvLLM実装を使用したエンドツーエンド推論サービングにおいて、SwiftKVは、最大2倍の集約スループットとトークンあたりの時間を60％短縮することができます。これにより、4x H100 GPU上で16ビット精度でLlama-3.1-70Bの場合、1秒あたり16,000トークンの推論スループットを実現し、GPUあたり560 TFlopsの驚異的な推論スループットを達成できます。

English

LLM inference for popular enterprise use cases, such as summarization, RAG, and code-generation, typically observes orders of magnitude longer prompt lengths than generation lengths. This characteristic leads to high cost of prefill and increased response latency. In this paper, we present SwiftKV, a novel model transformation and distillation procedure specifically designed to reduce the time and cost of processing prompt tokens while preserving high quality of generated tokens. SwiftKV combines three key mechanisms: i) SingleInputKV, which prefills later layers' KV cache using a much earlier layer's output, allowing prompt tokens to skip much of the model computation, ii) AcrossKV, which merges the KV caches of neighboring layers to reduce the memory footprint and support larger batch size for higher throughput, and iii) a knowledge-preserving distillation procedure that can adapt existing LLMs for SwiftKV with minimal accuracy impact and low compute and data requirement. For Llama-3.1-8B and 70B, SwiftKV reduces the compute requirement of prefill by 50% and the memory requirement of the KV cache by 62.5% while incurring minimum quality degradation across a wide range of tasks. In the end-to-end inference serving using an optimized vLLM implementation, SwiftKV realizes up to 2x higher aggregate throughput and 60% lower time per output token. It can achieve a staggering 560 TFlops/GPU of normalized inference throughput, which translates to 16K tokens/s for Llama-3.1-70B in 16-bit precision on 4x H100 GPUs.