MatKV: Trading Compute for Flash Storage in LLM Inference

TL;DR

MatKV tackles the rising cost and energy of LLM inference in retrieval-augmented generation by materializing key-value caches for retrieved documents on fast, inexpensive flash storage. By precomputing KV tensors during ingestion and loading them from SSD during inference, MatKV decouples prefill from decoding and enables overlap between I/O and computation, achieving up to 2× speedups and ~50% energy savings relative to GPU-based KV recomputation. The approach also enables decoding on low-end GPUs when KVs are preloaded, broadening deployment options while maintaining near-Vanilla QA accuracy across several benchmarks. With a principled ten-day rule for economics and observations of skewed document access, MatKV offers a practical, scalable storage-centric optimization for large-scale generative AI in RAG-heavy workloads.

Abstract

We observe two major trends in LLM-based generative AI: (1) inference is becoming the dominant factor in terms of cost and power consumption, surpassing training, and (2) retrieval augmented generation (RAG) is becoming prevalent. When processing long inputs in RAG, the prefill phase of computing the key-value vectors of input text is energy-intensive and time-consuming even with high-end GPUs. Thus, it is crucial to make the prefill phase in RAG inference efficient. To address this issue, we propose MatKV, a scheme that precomputes the key-value vectors (KVs) of RAG objects (e.g., documents), materializes them in inexpensive but fast and power-efficient flash storage, and reuses them at inference time instead of recomputing the KVs using costly and power-inefficient GPU. Experimental results using Hugging Face's Transformers library across state-of-the-art GPUs and flash memory SSDs confirm that, compared to full KV computation on GPUs, MatKV reduces both inference time and power consumption by half for RAG workloads, without severely impacting accuracy in the question-answering task. Furthermore, we demonstrate that MatKV enables additional optimizations in two ways. First, a GPU can decode text while simultaneously loading the materialized KVs for the next instance, reducing load latency. Second, since decoding speed is less sensitive to GPU performance than KV computation, low-end GPUs can be leveraged for decoding without significantly compromising speed once the materialized KVs are loaded into GPU memory. These findings underscore MatKV's potential to make large-scale generative AI applications more cost-effective, power-efficient, and accessible across a wider range of tasks and hardware environments.

Figures (10)

• Figure 1: GPU and SSD Cost/Performance Trend (2017 to 2024)
• Figure 2: Distribution of Accessed Vectors in RAG. While running 1M top-10 queries against a vector database with 9M document chunks, we measure the access frequency of each distinct vector. More than 900K documents are accessed twice or more. This skewed access indicates that there exists a large number of documents satisfying the ten-day rule.
• Figure 3: MatKV Architecture. (a) When a document (O) is inserted into the vector database, its KV cache is precomputed during the prefill phase using the LLM and stored in a commodity flash storage, enabling future reuse, (b) During inference, the system retrieves relevant documents based on the query (Q), loads their precomputed KV caches from flash storage, and directly utilizes them in the LLM, avoiding redundant prefill computations and improving inference efficiency.
• Figure 4: Overlapping in MatKV. In conventional architectures, all phases are strictly serialized, whereas MatKV enables concurrent KV loading and decoding.
• Figure 5: Prefill and decoding latency comparison between MatKV and Vanilla for a single request (LLaMA 3.1 70B). MatKV reduces prefill time to less than half of Vanilla, demonstrating the efficiency of flash storage over GPU recomputation.