SwiftKV: Inferencia Rápida Optimizada con Precarga con Transformación de Modelo Preservando el Conocimiento

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

October 4, 2024

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

cs.AI

Resumen

La inferencia de LLM para casos de uso empresariales populares, como la sumarización, RAG y la generación de código, generalmente observa longitudes de indicación órdenes de magnitud más largas que las longitudes de generación. Esta característica conlleva un alto costo de precarga y un aumento en la latencia de respuesta. En este documento, presentamos SwiftKV, un procedimiento novedoso de transformación y destilación de modelos diseñado específicamente para reducir el tiempo y el costo de procesamiento de tokens de indicación mientras se preserva la alta calidad de los tokens generados. SwiftKV combina tres mecanismos clave: i) SingleInputKV, que precarga la caché KV de capas posteriores utilizando la salida de una capa mucho más temprana, permitiendo que los tokens de indicación omitan gran parte de la computación del modelo, ii) AcrossKV, que fusiona las cachés KV de capas vecinas para reducir la huella de memoria y admitir un tamaño de lote más grande para un mayor rendimiento, y iii) un procedimiento de destilación que conserva el conocimiento y que puede adaptar LLMs existentes para SwiftKV con un impacto mínimo en la precisión y bajos requisitos de cómputo y datos. Para Llama-3.1-8B y 70B, SwiftKV reduce el requisito de cómputo de precarga en un 50% y el requisito de memoria de la caché KV en un 62.5%, con un deterioro mínimo en la calidad en una amplia gama de tareas. En el servicio de inferencia de extremo a extremo utilizando una implementación optimizada de vLLM, SwiftKV logra hasta 2 veces más rendimiento agregado y un 60% menos de tiempo por token de salida. Puede alcanzar un asombroso rendimiento de inferencia normalizado de 560 TFlops/GPU, lo que se traduce en 16K tokens/s para Llama-3.1-70B en precisión de 16 bits en 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.