SwiftKV : Inference rapide optimisée par préremplissage avec transformation de modèle préservant les connaissances

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

October 4, 2024

Auteurs: Aurick Qiao, Zhewei Yao, Samyam Rajbhandari, Yuxiong He

cs.AI

Résumé

L'inférence LLM pour des cas d'utilisation d'entreprise populaires tels que la résumé, RAG et la génération de code, observe généralement des longueurs de prompt de l'ordre de grandeur supérieures à celles des longueurs de génération. Cette caractéristique entraîne un coût élevé de préremplissage et une latence de réponse accrue. Dans cet article, nous présentons SwiftKV, une procédure de transformation et de distillation de modèle novatrice spécifiquement conçue pour réduire le temps et le coût de traitement des jetons de prompt tout en préservant une haute qualité des jetons générés. SwiftKV combine trois mécanismes clés : i) SingleInputKV, qui préremplit le cache KV des couches ultérieures en utilisant la sortie d'une couche bien plus tôt, permettant aux jetons de prompt de sauter une grande partie du calcul du modèle, ii) AcrossKV, qui fusionne les caches KV des couches voisines pour réduire l'empreinte mémoire et prendre en charge une taille de lot plus grande pour un débit plus élevé, et iii) une procédure de distillation préservant les connaissances qui peut adapter les LLM existants pour SwiftKV avec un impact minimal sur la précision et des exigences de calcul et de données faibles. Pour Llama-3.1-8B et 70B, SwiftKV réduit de 50% l'exigence de calcul du préremplissage et de 62,5% l'exigence mémoire du cache KV tout en entraînant une dégradation minimale de la qualité sur un large éventail de tâches. Dans le service d'inférence de bout en bout en utilisant une implémentation vLLM optimisée, SwiftKV réalise jusqu'à 2 fois plus de débit agrégé et 60% de temps inférieur par jeton de sortie. Il peut atteindre un impressionnant débit d'inférence normalisé de 560 TFlops/GPU, ce qui se traduit par 16K jetons/s pour Llama-3.1-70B en précision 16 bits sur 4x H100 GPUs.

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.