Compression method speeds up large language models on multiple GPUs
BasisKV: Value Latents for Communication and Key Routing
Artificial Intelligence
Summary
Longer conversations with large language models get slower because the computer has to remember more past information, which takes up a lot of memory and communication between GPUs. The authors found a new way to compress and share this past information efficiently when using multiple GPUs in parallel. Their method, called BasisKV, compresses the important parts and smartly guesses which data to fetch, making the whole process faster without much loss in quality. This approach led to up to a 42% speed increase in some setups.
What this means in practice
- •For machine learning engineers: Run long-context language model inference faster on multi-GPU servers by using improved key-value compression and communication reduction.
- •For cloud ai platform operators: Reduce GPU communication overhead during large model inference to serve longer queries more efficiently across distributed hardware.
Authors
Lixun Zhang, Zishan Shao, Kangning Cui, Yu Liang, Fan Yang, Yiran Chen
Abstract
Long-context large language model (LLM) inference is increasingly constrained by the cost of maintaining an ever-growing KV-cache. Low-rank K/V projections provide a promising way to compress the KV-cache with minimal quality loss, but prior work has largely optimized KV-cache compression in isolation, leaving distributed inference, especially with tensor parallelism where attention computation is partitioned across multiple GPUs, largely underexplored. We make two observations: (i) tensor-parallel communication is an important cost when deploying low-rank KV-cache compression in distributed inference; and (ii) Value compression can be associated with the output projection, allowing the same low-rank Value representation to reduce both Value-cache width and tensor-parallel communication. Based on these observations, we propose BasisKV, a low-rank KV-cache compression method for communication-intensive, long-context LLM inference. BasisKV jointly optimizes the low-rank factors across per-GPU attention partitions against the complete attention output, rather than fitting each partition in isolation. The compressed Value cache is further reused to help identify which Key pages should be accessed: it predicts part of the Key information needed for routing, while a Key-specific residual stores what cannot be recovered from Value. The original Keys are then retrieved only for the selected pages to compute exact QK scores. Across multiple LLM families and system settings, BasisKV achieves a better quality--efficiency tradeoff than existing low-rank baselines, reaching up to 1.42x end-to-end speedup under 8-GPU tensor-parallel inference. Our code is available at https://github.com/alexz949/BasisKV