SwiftKV: Inferência Rápida Otimizada com Preenchimento Antecipado por Transformação de Modelo Preservando Conhecimento

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

October 4, 2024

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

cs.AI

Resumo

A inferência LLM para casos de uso empresarial populares, como sumarização, RAG e geração de código, geralmente observa comprimentos de prompt ordens de magnitude maiores do que os comprimentos de geração. Essa característica resulta em alto custo de preenchimento e aumento da latência de resposta. Neste artigo, apresentamos o SwiftKV, um procedimento de transformação e destilação de modelo inovador, projetado especificamente para reduzir o tempo e o custo de processamento de tokens de prompt, ao mesmo tempo que preserva a alta qualidade dos tokens gerados. O SwiftKV combina três mecanismos-chave: i) SingleInputKV, que preenche o cache KV de camadas posteriores usando a saída de uma camada muito anterior, permitindo que os tokens de prompt pulem grande parte da computação do modelo, ii) AcrossKV, que mescla os caches KV de camadas vizinhas para reduzir a pegada de memória e suportar um tamanho de lote maior para maior rendimento, e iii) um procedimento de destilação que preserva o conhecimento e pode adaptar LLMs existentes para o SwiftKV com impacto mínimo na precisão e baixa exigência de computação e dados. Para o Llama-3.1-8B e 70B, o SwiftKV reduz o requisito de computação do preenchimento em 50% e o requisito de memória do cache KV em 62,5%, enquanto incorre em degradação mínima de qualidade em uma ampla gama de tarefas. No atendimento de inferência de ponta a ponta usando uma implementação vLLM otimizada, o SwiftKV alcança até 2 vezes mais rendimento agregado e 60% menos tempo por token de saída. Ele pode atingir um impressionante rendimento de inferência normalizado de 560 TFlops/GPU, o que se traduz em 16K tokens/s para Llama-3.1-70B em precisão de 16 bits em 4x GPUs 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.