PQCache: Product Quantization-based KVCache for Long Context LLM Inference

1 Introduction

• A 7B model processing 128K-length sequences with batch size 128 requires 1TB of KVCache, exceeding the 640GB GPU memory in an 8-card A100 setup.The resulting need to store cache in lower memory tiers adds transfer latency and scheduling complexity.Source: §1 · pp. 2–3
• Existing selective-attention methods either rely on assumptions that can hurt model quality or introduce notable latency.Dropping may discard tokens that gain importance later; offloading approaches can incur communication overhead or rely on block continuity that does not match discrete relevance.Source: §1 · p. 3
• Selective attention maps to information retrieval: prefilling supplies keys as embeddings for index construction, and decoding searches for relevant tokens using the new query.The paper frames KVCache storage and retrieval as an embedding retrieval problem.Source: §1 · pp. 3–4
• PQCache applies Product Quantization to keys during prefilling, then combines partitioned queries, centroids, and PQ codes during decoding to retrieve top-k key-value pairs.The method overlaps KVCache offloading, PQ construction, and code fetching with LLM computation, and uses a GPU block-level cache for frequently accessed pairs.Source: §1 · p. 4
• The authors present PQCache as the first work to apply information retrieval techniques to the KVCache memory challenge.They describe retrieval as an intuitive approach to managing large volumes of context information.Source: §1 · p. 4
• 4.60% higher InfiniteBench scores than existing methods, with low system latency, are reported for PQCache.The approach combines retrieval with overlapping and caching to minimize computation and communication overhead.Source: §1 · p. 4

2.1 Large Language Model Inference

• LLM inference uses causal self-attention, while FlashAttention reduces attention's space complexity to O(s) without changing its O(s^2) time complexity.Grouped-Query Attention uses fewer key/value heads than query heads.Source: §2.1 · p. 5
• During prefilling, the model computes and stores key-value pairs for input tokens; decoding computes the latest token and reuses the KVCache, whose memory grows linearly with sequence length.Each decoding step adds the newly generated token's key and value to the cache.Source: §2.1 · p. 6

2.2 Product Quantization

• Product Quantization splits each candidate embedding into m subspaces, clusters each into 2^b centroids, and stores m b-bit codes for approximate reconstruction and search.Search partitions the query, scores centroids, and aggregates scores through codes instead of comparing against every full embedding.Source: §2.2 · p. 6

2.3 GPU-CPU Memory Hierarchy

• The GPU–CPU hierarchy pairs faster, higher-bandwidth GPU memory with lower-bandwidth CPU memory connected through PCI-e or NVLink.The stated challenge is scheduling memory I/O alongside GPU computation to hide overhead; the design may extend to disk storage.Source: §2.3 · pp. 6–7

3 PQCache

• PQCache offloads the prefilling KVCache to CPU, compresses it with PQ, and fetches relevant key-value pairs on demand during decoding.Source: §3 · p. 7

3.1 Overview

• Attention scores generally follow power-law distributions, so PQCache retrieves only high-scoring tokens while retaining initial and recent local tokens for attention.The score distributions are shown for sampled positions in an XSUM example.Source: §3.1 · p. 7
• PQ partitions each key vector into m subspaces, clusters each separately, and stores per-token cluster codes using b bits per code.The resulting centroids have shape (m, 2^b, d_m), and codes have shape (s, m).Source: §3.1 · p. 8
• During decoding, query-centroid products and PQ codes estimate token scores; PQCache selects the top-k tokens and fetches their key-value pairs for attention.New tokens first remain local and receive PQ codes when evicted from the local window.Source: §3.1 · p. 8

3.2 Complexity Analysis

• Prefilling attention remains unchanged, while PQ construction adds sequence-length-linear work; decoding uses PQ scoring and selective attention to reduce the sequence-dependent computation.PQ structure communication is small enough to overlap with GPU computation, according to the complexity analysis.Source: §3.2 · p. 9

3.3 Prefilling Phase

• Prefilling overlaps GPU computation, asynchronous KVCache offloading, and CPU-side PQ construction to hide their added work.The scheduling design also targets PQ structure fetching and retrieval of selected tokens’ key-value pairs.Source: §3.3 · p. 9, §3.2 · p. 9

• The algorithm waits for PQ construction to finish in the next decoding phase.Source: §7 · p. 10
• The shown code lines connect completion of PQ construction with adding its data to the PQ collection.Source: §7 · p. 10
• The code appends the centroids and codes to PQ using PQ.append((Centroids, Codes)).Source: §7 · p. 10

• For sufficiently long sequences, GPU computation can fully cover communication and K-Means clustering, but clustering may outlast GPU work when sequences are shorter.The timing comparison covers one layer of a 7B LLM, PCI-e 1.0 (x16), and two Xeon Gold 6330 CPUs.Source: §17 · p. 11
• PQCache limits the maximum K-Means iterations adaptively so clustering can be overlapped by GPU computation.Source: §17 · p. 11
• The iteration cap is derived by matching modeled clustering time to GPU computation time, with clustering time fitted linearly in sequence length and iteration count.The maximum is clipped in practice to avoid too few or too many iterations.Source: §17 · p. 11
• The coefficients for the timing models are fitted by profiling single-layer computation and clustering across several sequence lengths for each model and device.Source: §17 · p. 11
• The shown decoding algorithm evicts a local token, assigns its key a PQ code, and asynchronously transfers its key and value to CPU.Source: §17 · p. 12

• PQ centroids remain on GPU and next-layer PQ codes can be prefetched during preceding-layer computation, allowing code communication to overlap decoding work.The centroids are fixed in size and independent of sequence length.Source: §3.4 · p. 12
• The decoding attention input combines initial tokens, selected top-k tokens, and local tokens.Middle-token key-value pairs are stored on CPU, while initial and local pairs remain on GPU.Source: §3.4 · p. 13
• A block-level GPU cache uses LFU or LRU eviction and checks for retrieved tokens before fetching missing key-value pairs from CPU.The only communication that cannot be overlapped is retrieval of top-k tokens, because it depends on the preceding PQ search.Source: §3.4 · p. 13

4.2 Model Performance

• With questions moved before the context, PQCache improves over SnapKV(C) and PyramidKV(C) by +7.10% across nearly all datasets.The authors report that those baselines decline when question placement varies, while PQCache does not rely on a question-position assumption.Source: Table 3 · p. 16, §4.2.1 · p. 16
• PQCache improves over baselines by +4.60% using 1/10 tokens on InfiniteBench and +3.90% using 1/10 tokens on LongBench.The reported gains are over baselines; the evaluations use 1/64 extra communication on InfiniteBench and 1/128 on LongBench.Source: Table 2 · p. 15, §4.2.1 · p. 15, Table 4 · p. 16
• Across InfiniteBench, PQCache performs best on most datasets and remains close to Full and Oracle, with an average difference from Oracle below 1.71%.A small number of individual tasks favor baselines, but the reported differences are not substantial.Source: Table 4 · p. 16, §4.2.2 · p. 17
• PQCache remains robust when combined with MInference, with only slight degradation relative to MInference alone despite shorter clustering and sparse attention affecting decoding.MInference accelerates prefilling, while PQCache addresses KVCache memory in decoding.Source: Table 5 · p. 18, §4.2.4 · p. 18
• On Llama-3.1-70B, PQCache is comparable to the uncompressed baseline even when CPU resources per GPU are halved.The authors attribute the small gap to larger-model robustness and more clustering iterations under the specified resource setup.Source: §4.2.5 · p. 18, Table 6 · p. 19
• On GSM8k CoT, PQCache outperforms the listed baselines across token counts; on the needle test, it retrieves the needle in most scenarios, unlike several baselines.H2O(C) performs better than PQCache at 1/10 tokens on GSM8k when allowed to access more tokens; the needle test extends to 131,000 tokens.Source: §4.2.3 · p. 17, Figure 9 · p. 17, §4.2.6 · p. 19
• PQCache is robust across tested PQ configurations except 8 × 2 on Qasper; 2 × 6 gives the best reported performance and is the default.More clustering iterations generally improve accuracy but increase latency; the adaptive strategy yields the least TT2T by avoiding GPU blocking.Source: §4.2.7 · p. 19, Figure 12 · p. 21, §4.3.4 · p. 22

4.3 Efficiency

• With overlapping and adaptive clustering, PQCache achieves nearly the lowest TT2T; its block-level cache cuts TPOT by 26.3% at 4K and 32.8% at 8K versus no cache.Prefilling offload time is negligible relative to GPU and KMeans computation, while overlapped decoding and GPU caching reduce end-to-end time and keep it stable as input length grows.Source: §4.3.1 · p. 20, Figure 11 · p. 21, §4.3.2 · p. 21

Hyper-parameter Tuning

• Practitioners should tune top-k to workload and GPU memory, then choose m × b to balance vector representation against clustering cost.The paper reports slight performance variation when m × b is held constant.Source: Hyper-parameter Tuning · p. 22

Multiple Requests and Multiple GPUs

• PQCache computations and communications extend to batched requests, but CPU-intensive workloads may block GPU computation as batch size increases.Adding CPUs is offered as a solution; the approach also applies to multiple-GPU scenarios.Source: Multiple Requests and Multiple GPUs · p. 22

Longer Output Sequences and Multiple-turn Conversation

• For longer outputs, input-based PQ structures may miss new information; periodically reconstructing PQ is a straightforward update.For multi-turn conversations, the paper offers either prefilling again with prior tokens or clustering each input separately and adding its output codes.Source: Longer Output Sequences and… · pp. 22–23

GPU and CPU Computational Capability

• PQCache can be constrained by CPU capability because clustering iterations are limited to avoid blocking GPU computation; the paper says a proper GPU-to-CPU ratio suits it best.Less powerful GPUs could permit more K-Means iterations during PQ construction, potentially improving performance.Source: GPU and CPU Computational Capability · p. 23

Combined with LLM Inference/Serving Systems

• PQCache is described as compatible with Orca, vLLM, and DistServe, though integration may require engineering effort.For prefill-decode disaggregation, it adds communication for PQ structures, which the paper says is much smaller than KVCache communication.Source: Combined with LLM Inference/Serving… · p. 23

Retrieval for LLM Inference

• The paper presents PQCache as an early use of information retrieval in LLM inference and identifies IVF and HNSW as possible alternatives to classic PQ.It suggests further research could explore retrieval techniques for LLM inference.Source: Retrieval for LLM Inference · p. 23

Selective Attention for KVCache

• Selective-attention methods either discard tokens or fetch relevant key-value pairs on demand; greedy eviction can lose information during later decoding.Examples include retaining initial and recent tokens, selecting by attention scores, and using an auxiliary model to choose tokens.Source: Selective Attention for KVCache · p. 23
• SPARQ and InfLLM offload KVCache to CPU and selectively fetch relevant key-value pairs for each attention computation.The paper places PQCache in this on-demand fetching category.Source: Selective Attention for KVCache · p. 24

KVCache Quantization

• KVCache quantization compresses the cache directly, while other techniques address quantization residuals or combine quantization with token importance.The cited examples include,, and, respectively.Source: KVCache Quantization · p. 24

KVCache Scheduling

• KVCache scheduling research optimizes storage and access across memory hierarchies, including communication scheduling and reuse across multi-turn conversations.The passage names FlexGen, AttentionScore, and KVCache streaming.Source: KVCache Scheduling · p. 24

Embedding Management

• Embedding management research covers compression, retrieval, and key-value storage, which the paper presents as a potential source of methods for LLMs.The passage cites embedding compression, retrieval, and key-value storage.Source: Embedding Management · p. 24

7 Conclusion

• PQCache is a system-algorithm co-designed method for effective and efficient long-context LLM inference.The method incorporates product quantization to reduce memory and computation burden.Source: §7 · p. 24
• PQ codes and centroids support approximate nearest-neighbor search for important tokens used in attention, while overlapping and caching keep overhead negligible.Source: §7 · p. 24
• 4.60% model-quality improvement on InfiniteBench over existing methods, with low system latency.Source: §7 · p. 24