DOCSEARCHbyAlgolia
Submit Docs

KVCache-Factory

repository·main·Indexed 23 days ago

https://github.com/zefan-cai/kvcache-factory

A unified playground for evaluating KV cache compression, retrieval, merging, and quantization methods to enable long-context LLM inference. It supports a wide range of methods including StreamingLLM, H2O, SnapKV, Quest, MiniCache, PyramidKV, and HeadInfer for lossless CPU offloading. The library includes tools for LongBench and Needle In A Haystack evaluations, as well as scripts for benchmarking decoding latency and peak memory usage.

Tokens
4.4K
Snippets
10
Records
22
Agent score
80%

What's inside KVCache-Factory

  1. How HeadInfer works (Lossless Offloading)

    main

    Unlike compression methods, HeadInfer is a lossless offloading method. It keeps the full KV cache but stores it in one slot per (layer, kv head) pair. It streams these slots between CPU and GPU, prefetching the next head's cache asynchronously while the current head computes.

    Key Characteristics:

    • Memory Efficiency: Only a constant number of head-caches occupy GPU memory regardless of context length.
    • Performance Trade-off: It is a memory-for-latency trade. Expect a decode-latency cost due to per-head cache traffic.
    • Requirements: Requires --attn_implementation flash_attention_2.
    • Usage: It can be used with any runner that accepts --method (e.g., LongBench).
    python3 run_longbench.py \
  --method HeadInfer \
  --model_path /path/to/Llama-3-8B-Instruct \
  --attn_implementation flash_attention_2 \
  --save_dir ./results_long_bench \
  --use_cache True

  2. Understand the MiniCache runtime storage model

    main

    The MiniCachePairStore is used for cross-layer compression (depth-dimension). It stores compressed pairs of adjacent layers keyed by (current_layer_idx, previous_layer_idx).

    Key operations available on the store:

    • has_pair: Checks if a pair exists.
    • restore_pair: Restores the compressed pair.
    • retention_indices: Returns the indices of retained tokens.
    • clear_pair: Clears a specific pair.
    • clear: Clears the entire store.

    Note: Same-layer pairs are rejected.

  3. Understand the trade-offs between query-head and KV-head cache granularity

    main

    In GQA models, the choice of KV-cache granularity impacts both memory usage and model accuracy:

    query_head (Default)

    • Mechanism: repeat_kv is called before compression. Each query head can have a distinct set of kept tokens.
    • Memory: Inefficient. The stored cache is num_key_value_groups times larger than necessary. For Meta-Llama-3-8B (32 query heads, 8 KV heads), the cache is ~4x larger than an efficient GQA cache.
    • Accuracy: Matches the original uncompressed model behavior exactly.

    kv_head

    • Mechanism: Compression happens before repeat_kv. A single set of tokens is selected per KV head, shared by its group of query heads. This requires reducing scores (via mean, max, or sum) before performing top-k selection.
    • Memory: Efficient. Reduces the compressed cache size by approximately the num_key_value_groups factor.
    • Accuracy: Introduces a semantic change. Because query heads in a group might rank different tokens, reducing their scores to pick a single set of tokens can slightly alter model output. Benchmarks show near-parity for most tasks, though some aggregation-breadth tasks may see slight decreases.

  4. Supported KV Cache methods and tools

    main

    KVCache-Factory provides a unified interface for various KV cache research directions:

    Compression/Retrieval Methods: PyramidKV, SnapKV, Quest, NACL, Scissorhands, MiniCache, H2O, StreamingLLM, CAM, L2Norm, AdaKV, HeadKV, ThinK.

    Specialized Tools/Selectors:

    • Quest style: query-aware page/token selector.
    • NACL style: proxy/random eviction selector.
    • Scissorhands style: persistence-of-importance selector.
    • MiniCache style: cross-layer SLERP merge/restore.
    • LOOK-M style: pivot merge.
    • KVMerger style: weighted merge.

  5. Run LongBench Evaluation

    main

    Evaluate models on the LongBench dataset using run_longbench.py. You can also use the helper script scripts/scripts_longBench/eval.sh.

    Common Arguments:

    • --method: Select the compression method (e.g., pyramidkv, snapkv, h2o, headinfer). Note: headinfer is lossless and requires --attn_implementation flash_attention_2.
    • --model_path: Local or Hugging Face model path.
    • --datasets: Comma-separated list (e.g., --datasets narrativeqa,qasper). Defaults to all 16 datasets.
    • --attn_implementation: flash_attention_2, sdpa, or eager. --method think requires eager.
    • --max_capacity_prompts: Target KV cache budget per layer.
    • --kv_cache_granularity: query_head (default) or kv_head (GQA-efficient). Supported for snapkv, pyramidkv, h2o, streamingllm, cam, l2norm, adakv, and headkv.
    • --quant_method: kivi, kvquant, or gear.
    • --nbits: Quantization bit width.
    • --eval_batch_size: Must be 1. Batching is not supported.
    export CUDA_VISIBLE_DEVICES=0

python3 run_longbench.py \
  --method pyramidkv \
  --model_path /path/to/Llama-3-8B-Instruct \
  --max_capacity_prompts 128 \
  --attn_implementation flash_attention_2 \
  --save_dir ./results_long_bench \
  --use_cache True

  7. Install KVCache-Factory

    main

    To install KVCache-Factory, clone the repository, install the base requirements, and add the current directory to your PYTHONPATH to ensure modules are discoverable.

    Note on FlashAttention: flash-attn is optional if you use --attn_implementation sdpa or eager, but it is required for FlashAttention v2 experiments. If needed, install it manually after torch using:

    pip install flash-attn --no-build-isolation

    Note on MInference: For the optional MInference integration (--method minference), install the specific requirements:

    pip install -r requirements-minference.txt
    git clone https://github.com/Zefan-Cai/KVCache-Factory.git
cd KVCache-Factory
pip install -r requirements.txt
export PYTHONPATH="$PWD:${PYTHONPATH}"

  8. Run Needle In A Haystack Evaluation

    main

    Run the Needle-in-a-haystack retrieval benchmark using run_needle_in_haystack.py. Supported methods for this runner are full, pyramidkv, snapkv, streamingllm, h2o, and cam.

    After inference, use scripts/scripts_needle/visualize.py to generate visualizations. You must update the FOLDER_PATH variable in the visualization script before running it.

    python -u run_needle_in_haystack.py \
  --s_len 1000 \
  --e_len 8001 \
  --model_provider LLaMA3 \
  --model_name /path/to/Llama-3-8B-Instruct \
  --attn_implementation flash_attention_2 \
  --step 100 \
  --method pyramidkv \
  --max_capacity_prompt 96 \
  --model_version Llama3_pyramidkv_96_test

  9. Benchmark Latency and Memory

    main

    To compare decoding latency and peak memory usage for a single prompt, use the scripts/benchmark_latency_memory.py script.

    python scripts/benchmark_latency_memory.py \
  --model_path /path/to/model \
  --method pyramidkv \
  --attn_implementation flash_attention_2 \
  --max_capacity_prompt 512 \
  --max_new_tokens 256 \
  --repeat 3

  11. Configure KV Cache quantization methods

    main

    KVCache-Factory supports several quantization methods via the --quant_method CLI flag.

    • kivi: Uses a KIVI-style configuration where the key is quantized per-channel and the value is quantized per-token.
    • kvquant: A partially implemented method for KVQuant-style outlier paths.
    • gear: A near-lossless compression recipe using uniform quantization and SVD-based low-rank residual reconstruction. When using gear, you must also provide --rank and --outlier_ratio.

  12. Configure GQA KV-Cache Granularity and Score Aggregation

    main

    When using Grouped Query Attention (GQA) models, you can control the granularity of the KV-cache layout and how scores are aggregated during compression using specific CLI flags.

    By default, the system uses query_head granularity, which stores the cache at the query-head level. This is memory-inefficient for GQA models (e.g., a 4x increase in size for Llama-3-8B) but preserves the original token-selection semantics. Switching to kv_head granularity allows for a more efficient cache layout but requires an aggregation strategy to decide which tokens to keep when multiple query heads share a single KV head.

    Supported flags:

    • --kv_cache_granularity: Choose between {query_head, kv_head}. Default is query_head.
    • --gqa_score_agg: When using kv_head granularity, choose how to reduce scores from the query heads in a group: {mean, max, sum}. Default is mean.