KVCache-Factory

https://github.com/zefan-cai/pyramidkv
Admin
KVCache-Factory is a unified framework for KV Cache compression in large language models, supporting...
Tokens:47,323
Snippets:866
Trust Score:8.6
License:MIT
Update:4 months ago
Show doc for...
Tokens:
Raw
### PyramidKV Installation

Source: https://github.com/zefan-cai/pyramidkv/blob/main/README.md

Provides commands to clone the PyramidKV repository and install it along with its dependencies.

```bash
git clone https://github.com/Zefan-Cai/PyramidKV.git
cd PyramidKV
pip install -r requirements.txt .
```

--------------------------------

### Lisp Conditional Expression Example

Source: https://github.com/zefan-cai/pyramidkv/blob/main/data/PaulGrahamEssays/diff.txt

Demonstrates how to use conditional expressions in Lisp, showing both if-then-else and assignment within conditionals.

```lisp
(if foo (= x 1) (= x 2))
```

```lisp
(= x (if foo 1 2))
```

--------------------------------

### Load Model, Tokenizer, and Input

Source: https://github.com/zefan-cai/pyramidkv/blob/main/examples/visualization.ipynb

Loads a pre-trained language model and tokenizer, and reads a sample input prompt from a file. Ensure the model path and input file are correctly specified.

```python
import os
import sys
sys.path.append(os.path.abspath("../"))
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# load the model and tokenizer
MODEL_NAME_OR_PATH = "meta-llama/Llama-3.2-3B-Instruct"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME_OR_PATH)

# load the sample input
with open("../pyramidkv/viztools/test_input.txt", 'r') as file:
    prompt = file.read().strip()
```

--------------------------------

### LongBench Inference Script Configuration

Source: https://github.com/zefan-cai/pyramidkv/blob/main/README.md

Configures and runs the LongBench inference script with various parameters like attention implementation and model path. This script is used to reproduce results presented in the paper.

```bash
export CUDA_VISIBLE_DEVICES=$1

method=$2 # Support PyramidKV, SnapKV, H2O, StreamingLLM
max_capacity_prompts=64 # 128,2048 in paper
attn_implementation=$3 # Support "flash_attention_2", "sdpa", "eager".
source_path=$4
model_path=$5
save_dir=${source_path}"results_long_bench" # path to result save_dir

python3 run_longbench.py \
    --method ${method} \
    --model_path ${model_path} \
    --max_capacity_prompts ${max_capacity_prompts} \
    --attn_implementation ${attn_implementation} \
    --save_dir ${save_dir} \
    --use_cache True
```
Tokens:
Raw