SwiftKV: Snelle Inference geoptimaliseerd voor vooraf invullen met Kennisbehoudende Modeltransformatie

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

October 4, 2024

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

cs.AI

Samenvatting

LLM-inferentie voor populaire zakelijke toepassingen, zoals samenvatting, RAG en codegeneratie, observeert doorgaans promptlengtes die qua orde van grootte veel langer zijn dan generatielengtes. Deze eigenschap leidt tot hoge kosten voor voorinvulling en verhoogde responstijd. In dit artikel presenteren we SwiftKV, een nieuw modeltransformatie- en distillatieproces dat specifiek is ontworpen om de tijd en kosten van het verwerken van prompttokens te verminderen, terwijl de hoge kwaliteit van gegenereerde tokens behouden blijft. SwiftKV combineert drie belangrijke mechanismen: i) SingleInputKV, dat de KV-cache van latere lagen voorvult met de uitvoer van een veel eerdere laag, waardoor prompttokens een groot deel van de modelberekening kunnen overslaan, ii) AcrossKV, dat de KV-caches van aangrenzende lagen samenvoegt om de geheugenfootprint te verkleinen en een grotere batchgrootte te ondersteunen voor een hogere doorvoer, en iii) een kennisbehoudende distillatieprocedure die bestaande LLM's kan aanpassen voor SwiftKV met minimale nauwkeurigheidsimpact en lage rekeneisen en gegevensvereisten. Voor Llama-3.1-8B en 70B vermindert SwiftKV de rekeneis van voorinvulling met 50% en de geheugenvereiste van de KV-cache met 62,5%, met minimale kwaliteitsdegradatie over een breed scala aan taken. In de end-to-end inferentieservice met behulp van een geoptimaliseerde vLLM-implementatie, realiseert SwiftKV tot 2x hogere totale doorvoer en 60% minder tijd per uitvoertoken. Het kan een verbazingwekkende genormaliseerde inferentiedoorvoer van 560 TFlops/GPU bereiken, wat zich vertaalt naar 16K tokens/s voor Llama-3.1-70B in 16-bits precisie op 4x H100 GPU's.

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.