SwiftKV: Быстрое инференцирование с оптимизацией предзаполнения с сохранением знаний при трансформации модели

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

October 4, 2024

Авторы: Aurick Qiao, Zhewei Yao, Samyam Rajbhandari, Yuxiong He

cs.AI

Аннотация

Инференция LLM для популярных корпоративных сценариев использования, таких как суммирование, RAG и генерация кода, обычно наблюдает порядки длиннее промптов, чем длины генерации. Эта характеристика приводит к высокой стоимости предзаполнения и увеличенной задержке ответа. В данной статье мы представляем SwiftKV, новую процедуру трансформации и дистилляции модели, специально разработанную для сокращения времени и стоимости обработки токенов промпта, сохраняя при этом высокое качество сгенерированных токенов. SwiftKV объединяет три ключевых механизма: i) SingleInputKV, который предзаполняет кэш KV более поздних слоев, используя вывод гораздо более раннего слоя, позволяя токенам промпта пропустить большую часть вычислений модели, ii) AcrossKV, который объединяет кэши KV соседних слоев для сокращения объема памяти и поддержки большего размера пакета для повышения пропускной способности, и iii) процедуру дистилляции, сохраняющую знания, которая может адаптировать существующие LLM для SwiftKV с минимальным влиянием на точность и низкими требованиями к вычислениям и данным. Для Llama-3.1-8B и 70B SwiftKV сокращает требования к вычислениям предзаполнения на 50% и требования к памяти кэша KV на 62.5%, сохраняя минимальное снижение качества на широком спектре задач. В инференции end-to-end с использованием оптимизированной реализации vLLM SwiftKV достигает до 2-кратного увеличения агрегированной пропускной способности и на 60% меньшего времени на вывод токена. Он может достичь потрясающих 560 TFlops/GPU нормализованной инференции, что переводится в 16K токенов/с для Llama-3.1-70B в 16-битной точности на 4x H100 GPU.

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.