SwiftKV: Fast Prefill-Optimized Inference with Knowledge-Preserving
  Model Transformation
Authors: name: Aurick Qiao; name: Zhewei Yao; name: Samyam Rajbhandari; name: Yuxiong He
Published: 2024-10-04 22:45:26

Abstract: 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.