DiffKV: Differentiated Memory Management for Large Language Models with Parallel KV Compaction

1 Introduction

• The KV cache grows with sequence length and concurrent requests, can exceed 90% of total memory, and limits serving efficiency.Its memory-bandwidth-bound nature also increases attention computation latency.Source: §1 · p. 1
• Existing pruning and quantization overlook differences in head sparsity, key and value roles, token significance, and information density across requests.The section frames these uniform policies as missing opportunities for finer-grained cache compression.Source: §1 · p. 1
• DiffKV differentiates key and value precision, token precision or pruning by importance, and critical-token allocation per head and request.Keys receive higher precision than values; tokens are stored at different precision levels or pruned according to attention-score importance.Source: §1 · p. 1
• In the illustrated requests, pruning and quantization apply uniform policies, while DiffKV adapts memory use across heads and requests.For Request 1, DiffKV uses 20.6% of memory averaged across both heads; Request 2 illustrates different allocations across heads and requests.Source: §1 · p. 2
• An on-GPU manager uses parallel KV compaction to pack fragmented free memory into contiguous regions across heterogeneous per-head, per-request allocations.The proposed design uses unified pages, a circular free-page list, and a bidirectional page table.Source: §1 · p. 2
• 2.7× to 5.7× KV cache compression and 1.9× to 5.4× throughput improvement are reported with near-lossless accuracy on complex reasoning workloads.The evaluation includes thinking models and long-generation capabilities.Source: §1 · p. 2

2.1 Large Language Models

• Transformers use causal attention over preceding tokens, and multi-head attention computes independently across heads while grouped-query attention shares keys and values across query-head groups.The KV cache stores earlier keys and values across generation steps; its size grows with sequence length and batch size.Source: §2.1 · p. 3

2.2 KV Cache Optimization

• PagedAttention allocates fixed-token pages on demand as sequences grow, reducing waste from reserving maximum-length cache space.The described approach enables larger batch sizes to improve serving efficiency.Source: §2.2 · p. 3
• Quantization reduces cache size with low-bit values, while pruning evicts less important tokens using attention scores; both can use uniform or static allocations.The reviewed pruning methods include uniform head budgets and layer-based heuristics that do not dynamically adapt to workloads.Source: §2.2 · p. 3
• CPU offloading selectively loads important tokens to the GPU but does not inherently reduce cache size and adds transfer latency.The section presents offloading as complementary to pruning.Source: §2.2 · p. 3

3.1 Differentiated Impacts of Keys and Values

• Keys influence attention scores for all tokens through the shared softmax denominator, while values affect only their own token’s contribution.In Llama3-8B, attention scores span seven orders of magnitude, compared with at most two for value-vector norms; K8V4 and K4V2 also outperform their mirror configurations.Source: §3.1 · p. 4

3.2 Differentiated Token Importance

• Token importance varies with attention scores, motivating high precision for the most important tokens, lower precision for moderately important ones, and pruning for the least important.This hierarchy offers a smoother trade-off between memory savings and generation quality than uniform quantization or pruning alone.Source: §3.2 · p. 4

3.3 Per-head Dynamic Sparsity Patterns

• The number of critical tokens needed to retain 95% of total attention varies considerably across layers and requests.Figure 4 reports layer averages aggregated across KV heads, with standard deviations capturing request variability.Source: §3.3 · p. 4, Figure 4 · p. 5
• Within a layer, critical-token counts vary across KV heads, and counts for an individual head also vary substantially across requests.Some heads show high request-level standard deviations.Source: §3.3 · pp. 4–5, Figure 5 · p. 5

3.4 Main Insights and Implications

• DiffKV combines differentiated key/value precision, hierarchical token compression, and per-request, per-head memory management to exploit these differences.Its contributions include request-aware precision mixes, sequence-length-adaptive token retention, and dynamic per-head allocation.Source: §3.4 · p. 5
• An on-GPU memory manager handles irregular per-head allocation patterns, translating KV-cache memory savings into performance gains.The paper presents this as a scalable memory-management contribution.Source: §3.4 · p. 5

KV compression policy

• DiffKV assigns keys higher precision than values and places tokens into high-precision, low-precision, or pruned tiers according to significance.The mix adapts to each request and sequence length, retaining more high-precision tokens for short sequences and compressing long sequences more aggressively.Source: §4 · p. 5

Prompt phase

• In the prompt phase, token significance averages attention scores received from subsequent tokens; the most recent W tokens are kept at high precision, typically with W set to 64.For GQA and MHA, scores from attention heads associated with a KV head are aggregated using the maximum operation.Source: §4 · p. 6

Generation phase

• During generation, the earliest token in the recent window becomes a compression candidate, and the least significant stored token may be downgraded or pruned.The policy re-quantizes an evicted high-precision token to low precision when warranted, pruning it only if it remains insignificant.Source: §4 · pp. 5–6

Discussion

• The policy uses two precision levels and shared thresholds across heads; per-head threshold tuning is left for future work.The authors report that shared thresholds suffice to capture differing head sparsity patterns in their empirical studies.Source: §4 · p. 6

5.1 Challenges

• Differentiated precisions make fixed-format pages wasteful and can misalign memory accesses, while per-head dynamic allocation raises compaction complexity to O(#requests × #heads) per inference step.Separate structures for high- and low-precision tokens would also increase metadata overhead.Source: §5.1 · p. 7

5.2 Parallel KV Compaction

• Parallel KV compaction performs per-head dynamic memory allocation and recycling on the GPU, using planning and coordination phases that can be parallelized.Planning is independent across heads; coordination can use a parallel prefix sum when the free-memory region is contiguous.Source: §5.2 · p. 7
• Unified pages store tokens at their allocated precision, while a circular free-page list and bidirectional page table support page management with reduced metadata overhead.The bidirectional table places high-precision page IDs at one end and low-precision IDs at the other, avoiding separate precision-specific lookups.Source: §5.2 · pp. 7–8

5.3 KV Compaction Workflow

• In the prompt workflow, heads plan their page needs and a prefix sum reclaims unused pages; during generation, a head allocates a page only when its current pages are full.Pages are recycled when a request finishes, and compaction runs once per inference step across the batch and model layers.Source: §5.3 · p. 8

6.1 DiffKV Architecture

• Each GPU worker pairs a memory manager for its assigned attention heads with an execution engine that compresses KV tensors and computes attention using the compressed cache.The scheduler batches requests within available GPU memory and sends them to the workers.Source: §6.1 · p. 9

6.2 Efficient Attention Kernel

• The attention kernel assigns each thread block one head per sequence and processes high-precision pages before low-precision pages to mitigate load imbalance.Each page first computes query-key dot products for attention scores, then a weighted sum of values.Source: §6.2 · p. 9
• Keys use a layout that makes warp-wide memory accesses contiguous, enabling coalescing while threads fetch vectorized elements and compute partial query-key dot products.The layout uses Kvec, the vectorization factor, and Kgroup, the number of threads per group.Source: §6.2 · p. 9
• Values are vectorized across tokens rather than features to reduce register pressure and support reduction across tokens.Thread groups reduce over the feature dimension, then a tree reduction combines their partial results.Source: §6.2 · p. 9
• For ultra-long sequences, the kernel splits sequence processing across thread blocks and merges the results with minimal computation overhead.Source: §6.2 · p. 9

7.1 Experiment Setup

• Experiments cover Llama3, Qwen2.5, and thinking models across general-knowledge, mathematics, code, long-context, and challenging reasoning benchmarks.Throughput and latency are measured on NVIDIA L40 GPUs, with model weights stored in FP16.Source: §7.1 · p. 10

7.2 Differentiated KV Compression Policy

• K8V4 matches FP16 accuracy across tested models and benchmarks, while mirror settings that assign lower precision to keys show accuracy losses, especially for Qwen2.5-7B.K8V2 retains at least 83% of FP16 accuracy across models; K4V1 yields almost zero accuracy, supporting the selected K8V4–K4V2 scheme for important and less-significant tokens.Source: §7.2 · p. 10
• Per-head dynamic sparsity outperforms static sparsity, while token pruning and precision allocation vary across benchmarks, with MMLU showing the highest sparsity and HumanEval+ the least.Dynamic sparsity maintains full Llama3-8B accuracy with 50% of GSM8K tokens pruned and with 10% of HumanEval+ tokens pruned.Source: §7.2 · pp. 11–12, Figure 12 · p. 12
• For evaluated non-thinking models, DiffKV achieves near-lossless accuracy relative to FP16, with an average degradation of 0.3% while using 19.3% to 36.7% memory.Source: §7.2 · p. 11
• 0.4% average accuracy degradation across LongBench categories is reported, with DiffKV using 14.9% memory for Llama3-8B and 27.0% for Qwen2.5-7B.Baseline methods are tuned to 25% memory use; DiffKV achieves superior accuracy across categories while using less or comparable memory.Source: §7.2 · p. 12, Table 2 · p. 13
• DiffKV achieves FP16-comparable generation quality on thinking models using 27.4%, 29.4%, and 23.5% of memory for QwQ-32B, R1-Distill-Qwen-14B, and R1-Distill-Llama-8B, respectively.On AIME24, 4-bit KV quantization degrades accuracy by 16.0%; on GPQA, pruning and partial-loading methods drop to almost zero accuracy at 50% of full KV memory.Source: §7.2 · p. 12, Table 3 · p. 13

7.3 System Performance

• On-GPU parallel compaction reduces memory-management latency by up to three orders of magnitude versus CPU multithreading, with overhead below 0.2% in prompt processing and 0.9% in generation.Model execution accounts for 96–97% of prompt-phase latency and 92–93% of generation-phase latency.Source: §7.3 · pp. 12–13, Figure 13 · p. 13, Figure 14 · p. 13
• 5.4× higher throughput over vLLM is reported for QwQ-32B on long MATH generations, alongside higher throughput than prior systems across evaluated models.DiffKV sustains a batch size of 15.9 for QwQ-32B, compared with vLLM’s 2.7.Source: §7.3 · pp. 13–14

• DiffKV combines differentiated key/value precision, hierarchical token-importance compression, and per-head dynamic sparsity.Source: §8 · p. 14
• The conclusion includes dynamic-workload comparisons of average latency, throughput, and achieved batch size.Source: Figure 16 · p. 14, Figure 17 · p. 14
• DiffKV compresses the KV cache 2.7×–5.7× with near-lossless accuracy and raises throughput 1.9×–5.4×, outperforming prior compression methods.Source: §8 · p. 14
• Parallel KV compaction handles irregular memory requirements across requests and attention heads, translating memory savings into performance gains.Source: §8 · p. 14