OjaKV

Official codebase for OjaKV: Context-Aware Online Low-Rank KV Cache Compression.

OjaKV targets the KV-cache memory bottleneck in long-context LLM inference. The repository contains:

• a calibration/compression pipeline for building low-rank checkpoints
• benchmark runners for LongBench and RULER
• implementations of the paper methods StaticPCA, OjaKV, and OjaKV-PF

Method Summary

StaticPCA

StaticPCA is the static low-rank baseline.

• It uses an offline low-rank basis computed from a calibration corpus.
• The basis is fixed at inference time.
• At inference, keys and values are reconstructed on the fly for compatibility with standard attention implementations.

In the paper, this corresponds to the strong static baseline that uses the same offline subspace but does not adapt online.

OjaKV

OjaKV is the main method proposed in the paper.

• It uses a hybrid KV storage policy: tokens with high reconstruction error are kept in full rank, while the rest are compressed.
• It updates the low-rank subspace online with Oja's algorithm.
• The basis is updated during both prefill and decoding, so it can track context shift.

OjaKV-PF

OjaKV-PF is the practical variant of OjaKV.

• It keeps the same compressed-cache logic as OjaKV.
• During prefill, attention is computed with the original full-rank keys and values.
• After prefilling, decoding uses the compressed cache.

PF stands for full-rank prefilling.

Repository Layout

OjaKV/
├── main_eigen_attn.py        # build low-rank checkpoints
├── eigen.sh                  # example compression script
├── models/                   # canonical model implementations
├── decompose/                # decomposition utilities
├── LongBench/
│   ├── run_longbench.py      # generate LongBench predictions
│   ├── eval.py               # score LongBench predictions
│   └── models/               # compatibility shims to repo-root models/
└── RULER/
    ├── run_ruler.py          # generate RULER predictions
    ├── eval_ruler.py         # score RULER predictions
    └── models/               # compatibility shims to repo-root models/

models/ at the repo root is the single source of truth. The LongBench/models and RULER/models directories are kept only as compatibility shims.

Environment

Recommended environment:

• Python 3.10
• CUDA-capable GPU
• PyTorch / FlashAttention versions compatible with requirements.txt

Setup:

conda create -n ojakv python=3.10
conda activate ojakv
pip install -r requirements.txt

If flash_attn needs to be rebuilt for your machine, install the matching PyTorch/CUDA stack first, then reinstall flash_attn.

1. Build Low-Rank Checkpoints

The compressed checkpoints consumed by StaticPCA, OjaKV, and OjaKV-PF are created by main_eigen_attn.py.

Quick Start

bash eigen.sh

eigen.sh now:

• treats the repository root correctly
• respects CUDA_VISIBLE_DEVICES if you set it
• writes checkpoints to saves/model/<net>/<error_rate>/

Direct Command

Example for Llama-3.1-8B-Instruct:

python main_eigen_attn.py \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --net Llama-3-1-8b \
  --cache_dir .cache/huggingface \
  --output_dir outputs/Llama-3-1-8b/0.075 \
  --save_dir saves/model/Llama-3-1-8b/0.075 \
  --calib_dataset wikitext2 \
  --nsamples 128 \
  --error_budget 0.075

The benchmark runners expect compressed checkpoints in:

saves/model/<net>/<error_rate>/

Supported net values in the current release:

• Llama-2-7b
• Llama-3-1-8b
• Llama-3-2-3b
• longchat-7b

2. Run LongBench

Predictions are written to:

LongBench/results/<model_name>/<dataset>/<method>.json
LongBench/results/<model_name>/<dataset>/<method>_<error_rate>.json

Full KV

python LongBench/run_longbench.py \
  --dataset narrativeqa \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --method FullKV \
  --eval_batch_size 1

StaticPCA

python LongBench/run_longbench.py \
  --dataset narrativeqa \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --net Llama-3-1-8b \
  --method StaticPCA \
  --error_rate 0.075 \
  --eval_batch_size 1

OjaKV

python LongBench/run_longbench.py \
  --dataset narrativeqa \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --net Llama-3-1-8b \
  --method OjaKV \
  --error_rate 0.075 \
  --eval_batch_size 1

OjaKV-PF

python LongBench/run_longbench.py \
  --dataset narrativeqa \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --net Llama-3-1-8b \
  --method OjaKV-PF \
  --error_rate 0.075 \
  --eval_batch_size 1

Evaluate LongBench

python LongBench/eval.py \
  --results_dir LongBench/results/Llama-3.1-8B-Instruct

The evaluation script automatically scans every *.json prediction file under the dataset directories and writes:

• LongBench/results/<model_name>/metrics.json
• LongBench/results/<model_name>/results.csv

3. Run RULER

Predictions are written to:

RULER/results/ruler/<model_name>/<context_length>/<dataset>/<method>.json
RULER/results/ruler/<model_name>/<context_length>/<dataset>/<method>_<error_rate>.json

Example:

python RULER/run_ruler.py \
  --context_length 16384 \
  --dataset niah_single_1 \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --method FullKV \
  --eval_batch_size 1

Compressed variants:

python RULER/run_ruler.py \
  --context_length 16384 \
  --dataset niah_single_1 \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --net Llama-3-1-8b \
  --method StaticPCA \
  --error_rate 0.075 \
  --eval_batch_size 1

python RULER/run_ruler.py \
  --context_length 16384 \
  --dataset niah_single_1 \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --net Llama-3-1-8b \
  --method OjaKV \
  --error_rate 0.075 \
  --eval_batch_size 1

python RULER/run_ruler.py \
  --context_length 16384 \
  --dataset niah_single_1 \
  --model_path meta-llama/Llama-3.1-8B-Instruct \
  --net Llama-3-1-8b \
  --method OjaKV-PF \
  --error_rate 0.075 \
  --eval_batch_size 1

Evaluate RULER

python RULER/eval_ruler.py \
  --results_dir RULER/results/ruler/Llama-3.1-8B-Instruct/16384

The evaluation script automatically scans every *.json prediction file under the dataset directories and writes:

• RULER/results/ruler/<model_name>/<context_length>/metrics.json
• RULER/results/ruler/<model_name>/<context_length>/results.csv

Notes on Current Release

• The benchmark runners are aligned with the paper names: FullKV, StaticPCA, OjaKV, and OjaKV-PF.
• The LongBench and RULER runners now import the canonical implementations from the repo-root models/ directory.
• The current public release only exposes the Llama-family / LongChat path used in the paper. Legacy OPT/MPT compression branches have been removed from the public main flow.
• The runners truncate overlength inputs by keeping the head and tail of the prompt, matching the original repository behavior.

Citation

If you use this repository, please cite the accompanying OjaKV paper.