PiKV: KV Cache Management System for Mixture of Experts

TL;DR

PiKV is an open-source distributed KV cache system designed to optimize memory and communication efficiency for Mixture of Experts models in large-scale language model inference.

Contribution

PiKV introduces expert-sharded storage, routing, scheduling, and compression techniques to improve KV cache management in MoE architectures.

Findings

Reduces memory usage through compression modules.

Improves cache access efficiency with expert sharding and routing.

Open-source implementation available at GitHub.

Abstract

As large-scale language models continue to scale up in both size and context length, the memory and communication cost of key-value (KV) cache storage has become a major bottleneck in multi-GPU and multi-node inference. While MoE-based architectures sparsify computation across experts, the corresponding KV caches remain dense and globally synchronized, resulting in significant overhead. We introduce \textbf{PiKV}, a parallel and distributed KV cache serving framework tailored for MoE architecture. PiKV leverages \textit{expert-sharded KV storage} to partition caches across GPUs, \textit{PiKV routing} to reduce token-to-KV access, and a \textit{PiKV Scheduling} to adaptively retain query-relevant entries. To further reduce memory usage, PiKV integrates \textit{PiKV Compression} modules the caching pipeline for acceleration. PiKV is recently publicly available as an open-source software library: \href{https://github.com/NoakLiu/PiKV}{https://github.com/NoakLiu/PiKV}. PiKV is still a living project, aiming to become a comprehesive KV Cache management system for MoE Architectures.