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

要約

要約、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% 短縮を実現します。
560 TFlops/GPU という驚異的な正規化推論スループットを達成できます。これは、4x H100 GPU 上の 16 ビット精度で Llama-3.1-70B の場合 16K トークン/秒に相当します。
私たちのトレーニング、推論、モデルの実装はオープンソースであり、https://huggingface.co/collections/Snowflake/swiftkv-models-674f7d7474eb789e185d31cb から参照できます。

要約(オリジナル)

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. Our training, inference, and model implementations are open-sourced and can be found through https://huggingface.co/collections/Snowflake/swiftkv-models-674f7d7474eb789e185d31cb.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google