Zefan-Cai/KVCache-Factory

Unified KV Cache Compression Methods for Auto-Regressive Models

What it solves

KVCache-Factory addresses the memory bottleneck associated with the Key-Value (KV) cache during long-context Large Language Model (LLM) inference. As context length increases, the KV cache grows linearly, consuming vast amounts of GPU memory and slowing down decoding. This project provides a unified framework to implement, compare, and evaluate various strategies to reduce this memory footprint without sacrificing model performance.

How it works

The project acts as a "playground" that integrates multiple KV cache optimization techniques under a single evaluation interface. It supports several categories of methods:

  • Compression and Eviction: Removing less important tokens (e.g., StreamingLLM, H2O, SnapKV, PyramidKV).
  • Retrieval: Selecting only the most relevant tokens for the current query (e.g., Quest, L2Norm).
  • Merging: Aggregating multiple KV pairs into one (e.g., CAM).
  • Quantization: Reducing the precision of the cached values (e.g., KIVI, KVQuant, GEAR).
  • Offloading: Moving the KV cache to CPU memory and streaming it back to the GPU head-by-head to maintain full precision without GPU memory overflow (HeadInfer).

It includes runners for standard benchmarks like LongBench, RULER, and "Needle-in-a-haystack" to measure the impact of these methods on accuracy and latency.

Who it’s for

This tool is designed for AI researchers and developers working on LLM efficiency, specifically those looking to optimize long-context window performance or benchmark different KV cache management strategies.

Highlights

  • Unified Interface: Supports a wide array of baselines including PyramidKV, SnapKV, H2O, and StreamingLLM.
  • Diverse Optimization Strategies: Combines eviction, retrieval, merging, quantization, and lossless offloading in one place.
  • Comprehensive Benchmarking: Integrated support for LongBench, RULER, and Needle-in-a-haystack evaluations.
  • Hardware Flexibility: Supports multi-GPU inference for large models (e.g., Llama-3-70B) and provides paths for both FlashAttention v2 and SDPA.
  • Visualization Tools: Includes utilities to visualize attention patterns to understand how different compression methods affect model focus.

Related

  • Project
    NVIDIA/kvpress

    A framework for implementing and benchmarking KV cache compression methods to reduce memory usage and increase inference speed for long-context LLMs.

  • Dispatch
    NVIDIA KVPress: Toolkit for KV Cache Compression in Long-Context LLMs

    NVIDIA has released KVPress, a Python toolkit that implements state-of-the-art KV cache compression techniques to reduce the memory footprint and increase decoding speed of long-context Large Language Models.

  • Project
    ModelEngine-Group/unified-cache-management

    A cache management framework for LLMs that reduces inference latency by persisting the KV cache and using sparse attention retrieval to minimize redundant computations.

  • Project
    quantumaikr/quant.cpp

    A high-performance C library for LLM inference that uses KV cache compression to enable long-context document understanding on consumer hardware, reducing the need for chunk-based RAG.

  • Project
    LMCache/LMCache

    A KV cache management layer for LLM inference that reduces TTFT and improves throughput by making KV caches persistent and reusable across requests and engines.