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

TL;DR

SwiftKV is a novel model transformation technique that reduces prefill computation in large language models, significantly improving inference speed while maintaining high output quality, suitable for enterprise applications.

Contribution

It introduces a cache reuse method and knowledge-preserving distillation to optimize LLM inference, achieving 25-50% prefill FLOPs reduction with minimal accuracy loss.

Findings

Reduces prefill computation by 25-50% across LLMs

Doubles inference throughput and reduces token generation time by 60%

Achieves 560 TFlops/GPU inference throughput for Llama-3.1-70B

Abstract

LLM inference for enterprise applications, such as summarization, RAG, and code-generation, typically observe much longer prompt than generations, leading to high prefill cost and response latency. We present SwiftKV, a novel model transformation and distillation procedure targeted at reducing the prefill compute (in FLOPs) of prompt tokens while preserving high generation quality. First, SwiftKV prefills later layers' KV cache using an earlier layer's output, allowing prompt tokens to skip those later layers. Second, SwiftKV employs a lightweight knowledge-preserving distillation procedure that can adapt existing LLMs with minimal accuracy impact. Third, SwiftKV can naturally incorporate KV cache compression to improve inference performance in low-memory scenarios. Our comprehensive experiments show that SwiftKV can effectively reduce prefill computation by 25-50% across several LLM families while incurring minimum quality degradation. In the end-to-end inference serving, 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. SwiftKV is open-sourced at https://github.com/snowflakedb/arctictraining.