SwiftKV: Inferenza Veloce Ottimizzata per il Prefill con Trasformazione del Modello Preservante la Conoscenza

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

October 4, 2024

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

cs.AI

Abstract

L'inferenza LLM per casi d'uso aziendali popolari, come la sintesi, RAG e la generazione di codice, di solito osserva lunghezze dell'input dell'ordine di grandezza superiori rispetto alle lunghezze di generazione. Questa caratteristica porta a un elevato costo di prefill e a un aumento della latenza di risposta. In questo articolo, presentiamo SwiftKV, una procedura di trasformazione e distillazione del modello innovativa progettata specificamente per ridurre il tempo e il costo di elaborazione dei token di input preservando nel contempo l'alta qualità dei token generati. SwiftKV combina tre meccanismi chiave: i) SingleInputKV, che precompila la cache KV dei livelli successivi utilizzando l'output di un livello molto precedente, consentendo ai token di input di saltare gran parte del calcolo del modello, ii) AcrossKV, che unisce le cache KV dei livelli adiacenti per ridurre l'occupazione di memoria e supportare batch size più grandi per una maggiore capacità di elaborazione, e iii) una procedura di distillazione che conserva la conoscenza e può adattare i LLM esistenti per SwiftKV con un impatto minimo sull'accuratezza e un basso requisito di calcolo e dati. Per Llama-3.1-8B e 70B, SwiftKV riduce il requisito di calcolo del prefill del 50% e il requisito di memoria della cache KV del 62,5% pur causando una degradazione minima della qualità in una vasta gamma di compiti. Nell'inferenza end-to-end con un'implementazione vLLM ottimizzata, SwiftKV realizza un throughput aggregato fino a 2 volte superiore e un tempo per token di output inferiore del 60%. Può raggiungere un sorprendente throughput di inferenza normalizzato di 560 TFlops/GPU, che si traduce in 16K token/s per Llama-3.1-70B in precisione a 16 bit su 4 GPU H100.

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.