SwiftKV: Schnelle Inferenz mit optimierter Vorabfüllung und wissenserhaltender Modelltransformation

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

October 4, 2024

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

cs.AI

Zusammenfassung

LLM-Inferenz für beliebte Unternehmensanwendungsfälle wie Zusammenfassung, RAG und Code-Generierung beobachtet in der Regel Größenordnungen längere Eingabesequenzen als Generierungslängen. Diese Eigenschaft führt zu hohen Kosten für das Vorausfüllen und erhöhter Antwortlatenz. In diesem Papier präsentieren wir SwiftKV, ein neuartiges Modell-Transformations- und Destillationsverfahren, das speziell entwickelt wurde, um die Zeit- und Kostenbelastung bei der Verarbeitung von Eingabesequenzen zu reduzieren, während gleichzeitig die hohe Qualität der generierten Sequenzen erhalten bleibt. SwiftKV kombiniert drei Schlüsselmechanismen: i) SingleInputKV, das die KV-Cache späterer Schichten mit der Ausgabe einer viel früheren Schicht vorbefüllt, wodurch die Eingabesequenzen einen Großteil der Modellberechnung überspringen können, ii) AcrossKV, das die KV-Caches benachbarter Schichten fusioniert, um den Speicherbedarf zu reduzieren und eine größere Batch-Größe für eine höhere Durchsatzrate zu unterstützen, und iii) ein wissensbewahrendes Destillationsverfahren, das vorhandene LLMs für SwiftKV mit minimalem Genauigkeitsverlust und geringem Rechen- und Datenbedarf anpassen kann. Für Llama-3.1-8B und 70B reduziert SwiftKV den Rechenbedarf des Vorausfüllens um 50 % und den Speicherbedarf des KV-Caches um 62,5 %, wobei nur minimale Qualitätsverschlechterungen über eine breite Palette von Aufgaben hinweg auftreten. Bei der End-to-End-Inferenzbereitstellung mit einer optimierten vLLM-Implementierung erzielt SwiftKV eine um bis zu 2-fach höhere Gesamtdurchsatzrate und eine um 60 % niedrigere Zeit pro Ausgabesequenz. Es kann eine beeindruckende normierte Inferenzdurchsatzrate von 560 TFlops/GPU erreichen, was sich in 16K Token/s für Llama-3.1-70B in 16-Bit-Präzision auf 4x H100-GPUs übersetzt.

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.