KVPR : Efficient LLM Inference with I / O -Aware KV Cache Partial Recomputation (2025)

Abstract: Inference for Large Language Models (LLMs) is computationally demanding. To reduce the cost of auto-regressive decoding, Key-Value (KV) cache is used to store intermediate activations, which significantly lowers the computational overhead for token generation. However, the memory required for the KV cache grows rapidly, often exceeding the capacity of GPU memory. A cost-effective alternative is to offload KV cache to CPU memory, which alleviates GPU memory pressure, but shifts the bottleneck to the limited bandwidth of the PCIe connection between the CPU and GPU. Existing methods attempt to address these issues by overlapping GPU computation with I/O or employing CPU-GPU heterogeneous execution, but they are hindered by excessive data movement and dependence on CPU capabilities. Fully overlapping PCIe communication latency gets challenging as the size of the KV cache grows and/or the GPU compute capabilities increase. In this paper, we introduce KVPR, an efficient I/O-aware LLM inference method where the CPU first transfers a partial set of activations, from which the GPU can start recomputing the KV cache values. While the GPU recomputes the partial KV cache, the remaining portion of the KV cache is transferred concurrently from the CPU. This approach overlaps GPU recomputation with KV cache transfer to minimize idle GPU time and maximize inference performance. KVPR is fully automated by integrating a profiler module that utilizes input characteristics and system hardware information, a scheduler module to optimize the distribution of computation and communication workloads, and a runtime module to efficiently execute the derived execution plan. Experimental results show that KVPR achieves up to 35.8% lower latency and 46.2% higher throughput during decoding compared to state-of-the-art approaches. The code is available at https://github.com/chaoyij/KVPR.

## introduction large language models (llms) have made remarkable progress in recent years, demonstrating their ability to power diverse applications such as machine translation @xcite , summarization @xcite , creative content generation @xcite , and personalized recommendations @xcite . real-time applications, including conversational agents and live translation @xcite , depend on low latency to provide seamless user interaction, while large-scale deployments require high throughput to support concurrent users and process substantial data efficiently @xcite . key-value (kv) cache is essential in autoregressive decoding for llms, as it stores the intermediate key and value activations from earlier steps in the attention mechanism. this reduces the computational complexity of generating each token from quadratic to linear by eliminating the need to recompute these activations for every generated token. however, this comes at a cost: the size of the kv cache grows linearly with batch size, sequence length, and model size, leading to substantial memory demands @xcite . gpu memory, while optimized for highbandwidth access by computation units, is inherently limited and often insufficient to handle the large and growing size of the kv cache. one cost-effective approach to address this limitation is to offload the kv cache to cheaper and plentiful cpu memory, and could be further offloaded to hard disks and network storage @xcite . while offloading reduces gpu memory pressure, it introduces a new bottleneck: the slow pcie bus becomes a limiting factor when transferring the kv cache from cpu to gpu for computation. due to the long pcie transfer time, the overall decoding latency increases and token generation throughput decreases, hindering the overall inference efficiency of the system @xcite . to evaluate the impact of communication overhead, we set up an llm inference serving system (shown in figure 1 ) using an nvidia a100 gpu. data transfer between the cpu dram and gpu hbm occurs over a pcie 4.0 16 lanes with a bandwidth of 32 gb/s. to mitigate the issue of low latency and bandwidth of pcie, flexgen @xcite and pipeswitch @xcite attempt to overlap gpu computation of the current layer with kv cache loading for the next layer. however, the effectiveness of such an overlap is capped by the task that takes the longest time. in most systems, pcie transfer time overshadows gpu computation latency, particularly with large batch and context sizes. hence, fully overlapping gpu computation with pcie transfer time is infeasible. fastdecode @xcite suggests computing attention scores directly on the cpu, which has faster access to the kv cache compared to the gpu. similarly, hetegen @xcite , twinpilots @xcite , and park and egger employ cpu-gpu heterogeneous execution to hide data transfer overhead by performing computations on the cpu. however, as demonstrated later in our results, such an approach puts a burden on the cpu to satisfy the kv cache computation demands from multiple gpus attached to a cpu host, thereby limiting scalability. in this paper, we propose kvpr, a novel approach for efficient llm inference that balances the gpu computation and pcie bandwidth tradeoffs. instead of transferring the entire kv cache from cpu to gpu to compute an attention score, the cpu transfers a partial set of activations, which are smaller in size and are required to generate part of the kv cache, to the gpu. the gpu then starts recomputing the partial kv cache from the input activations. concurrently, the cpu transfers the remaining kv cache over pcie. kvpr ensures the computation of exact attention scores without approximation, while minimizing gpu idle time and improving overall latency and throughput. kvpr achieves a near-perfect overlap of pcie transfer time and gpu recomputation time by determining the optimal fraction of activations that need to be recomputed. kvpr is fully automated in determining the recomputation and communication split. it includes a profiler module that collects system hardware information, a scheduler module that formulates the problem as a linear programming problem to determine the optimal split point, and a runtime module that manages memory allocation on both devices and coordinates data transfer between them. experimental results show significant improvements in inference latency or throughput, depending on workload. in summary, our contributions are as follows: in the decoding stage, the i-th decoder layer receives a single token x @xmath0. the kv cache is updated by concatenating the newly computed key and value pairs with the existing ones: the remaining attention and feedforward computations in the decoding stage are identical to those in the prefilling stage. column-by-column scheduling (appendix a.1) is more effective for maximizing throughput by increasing the effective batch size (number of batches times batch size) to process more sequences in parallel, at the cost of longer latency. in this approach, the model weights are offloaded to cpu memory to accommodate a large batch size. the model weights and kv cache for a single layer are transferred to gpu memory and processed for the first batch. instead of moving to the next layer for the current batch, subsequent batches are processed using the same layer while keeping the weights stationary in gpu memory. once a group of batches is processed for the first layer, the process moves to the second layer for each batch. note that the effective batch size is limited by the available storage for activations and kv cache, as they must still be stored in cpu memory or external storage once they exceed the gpu memory capacity. our proposed design is independent of the scheduling strategy, whether row-by-row or column-by-column, and aims to overlap the majority of the pcie transfer time with gpu computations, thereby improving overall efficiency. ## related works to address the memory demands of llms in resource-constrained settings, offloading techniques aim to minimize the latency of data transfer between cpus and gpus. flexgen @xcite proposes to offload weights, activations, and kv cache to cpu memory or external storage and maximizes throughput for larger batch sizes by formulating the optimization as a graph traversal problem. hetegen @xcite uses the cpu for partial computation on offloaded weights while transferring the remaining workload to the gpu. twinpilots @xcite further optimizes workload balancing between the cpu and gpu at the operator level. fastdecode @xcite reduces kv cache data movement by offloading the kv cache and attention computation entirely to the cpu. @xcite overlap gpu linear projection computations with cpu-based attention computations across multiple batches to improve resource utilization. alisa @xcite compresses the kv cache based on sparsity and offloads kv cache exceeding gpu memory capacity. when loading the kv cache to the gpu, alisa recomputes a portion of the kv cache first and then transfers the remainder, where we propose overlapping the recomputation and transfer by adaptively determining the optimal split point. furthermore, alisa addresses only the row-by-row schedule, while kvpr extends to the column-by-column schedule. kvpr is orthogonal to cpu-assisted and kv cache compression approaches, making it compatible for integration with these techniques to further improve overall system performance. as shown in the additional experiments provided in appendix a.7, we demonstrate that the cpu can become a bottleneck in certain distributed system configurations. in contrast, kvpr optimizes gpu utilization and data transfer efficiency without relying on additional cpu resources or approximations of the kv cache. ## conclusion in this paper, we introduce kvpr, an efficient cpu-gpu i/o-aware llm inference method designed to accelerate kv cache loading. kvpr minimizes the data transfer between the cpu and gpu by leveraging kv cache partial recomputation. by overlapping this recomputation with data transmission, kvpr significantly reduces idle gpu time and enhances overall inference performance. future work could extend our method to tolerate kv cache loading from remote network storage or scale to large multi-gpu infrastructure, further enhancing its applicability and performance in diverse deployment scenarios.