Skip to content

Repository files navigation

PriorityKV

Agent traces make schema fields and pleasantries equally expensive to cache but very unequal to lose. PriorityKV uses visible trace structure as a bounded eviction prior, then evaluates INT4 quality and packed-system cost as separate questions.

Arush Sharma — IIT (ISM) Dhanbad · Anupam Rawat — IIT Bombay
Apache-2.0 · Python 3.11–3.12 · Primary evaluation: Qwen3-8B on NVIDIA H200

Headline Result and claim boundary
Operating boundary (18,240 generations) Structural retention is at ceiling while the protected set fits the keep budget and collapses past oversubscription ratio ≈ 1.1. Replicated on Qwen3-8B, Llama-3.1-8B and Qwen3-32B — the 32B curve matches 8B to three significant figures at every level, so the boundary is a property of the policy, not the model.
Mechanism (per task family) State in the leading system block survives 18× oversubscription (tool-schema 1.000); mid-conversation state falls to 0.000. The ratio alone does not predict what is lost — the tiebreak position does.
Past the boundary Structure loses to matched-budget SnapKV in 20/24 Llama cells with b=0 (zero wins in the opposite direction). The 25% parity below is an operating-point property, not a general result.
ADAPT Never loses to SnapKV across 48 cells / 2 models; significantly exceeds it on Qwen at 2% keep (up to p=4.9e-4, c=0). Does not replicate on Llama, where SnapKV never degrades.
Qwen eviction, 25% keep (n=120) Structure 0.933 (112/120) vs uniform/random 0.008. Note: this baseline is position-blind on a benchmark whose gold sits outside the recent window, so treat it as motivation, not evidence.
Matched attention selectors (n=120) Structure 0.933 vs SnapKV/PyramidKV/hybrid 0.900; McNemar p=0.125, not significant.
Matched INT4 placement (n=240) FullKV 0.8875, uniform 0.8792, structure 0.8833: role-aware INT4 does not separate quality.
Packed H200 path Payload 0.719× and peak 0.868×, but E2E 1.11–1.12× and TPOT 1.20–1.21×: fewer bytes, higher latency.

External test — pre-registered prediction confirmed. The prediction in docs/PREREG_BFCL_TIGHT.md was committed before any tight-budget conversation was scored: the mechanism implies ADAPT should not help on BFCL, whose failures are losses of accumulated multi-turn state rather than of schemas. Measured at the rule-selected budget (keep_frac=0.15), n=233 paired, 0 budget violations:

Arm Accuracy vs SnapKV
FullKV 0.223
SnapKV 0.137
ADAPT 0.112 Δ=−0.026, p=0.327 (n.s.)

The method confers no external benefit; the mechanism that predicts its failure does have external predictive power. A pilot ladder also showed SnapKV floors below a 10% budget on BFCL, so the regime where the sweep shows an ADAPT advantage does not exist on that workload.

Paper: compiled PDF · LaTeX source · PDF build instructions

Qwen matched-budget eviction and attention baselines

Why structure?

Agent traces place tool schemas, superseding instructions, persistent identifiers, tool results, ordinary dialogue, and filler in one KV cache. PriorityKV asks whether the application-visible role of a token should influence retention when the cache must shrink. It tests eviction reliability separately from mixed-precision quality and systems cost.

Conceptual agent-trace failure mode

What is and is not claimed

The public claim boundary below is copied from docs/EVIDENCE.md, the canonical claim registry:

On PriorityBench-A (synthetic agent traces), structure-aware retention at a 25% keep budget far exceeds position-blind eviction on Qwen (P0 n=120: structure 0.933 vs uniform/random ~0.008). Mid-context relocation ties FullKV on s0/s1/s2 (e.g. s0 both 0.975); burying state hurts structure on all three slices (s0 0.675 vs FullKV 0.900; s1 0.650 vs 0.875; s2 0.675 vs 0.900) while uniform/random stay ~0 — so we do not claim structure beats FullKV. A CPU gold-span audit shows gold is not concentrated in sink+recent on Qwen or Llama (≈0–1% of gold tokens); structure keep retains essentially all gold tokens while uniform retains ≈0–1%, so the Qwen blind-eviction gap is retention-real, not a labeling leak into the always-kept window. Versus SnapKV-class selection on Qwen (n=120): structure 0.933 vs SnapKV/Pyramid/hybrid 0.900 (112/120 vs 108/120; McNemar p=0.125, not significant) — we claim only that it matches or slightly exceeds SnapKV-class methods while decisively beating position-only baselines. Hybrid did not beat SnapKV. Llama-3.1 at kf=0.25 is saturated among structure+attention arms (all 1.0); the gold audit rules out “gold already in sink+recent” as the explanation — the task is too easy once any competent keep runs. At kf=0.05 SnapKV outperforms structure on two slices (s0: 1.0 vs 0.875; s1: 1.0 vs 0.900). P2 streamed-cold is a smoke test (~36 GiB peak in log), not a systems result.

Not claimed:

  • Structure beats FullKV or SnapKV in general
  • Significant structure≫SnapKV on Qwen
  • Universal cross-model transfer to Llama
  • Soft INT4 quality win; peak VRAM collapse; LongBench/RULER matrices

The additional frozen-core negative remains explicit: 75% INT4 placement does not create a PriorityBench quality separation, and the current FlashInfer shim regresses TPOT.

System at a glance

PriorityKV maps chat messages to token roles, forms 16-token pages, applies a matched retention or precision budget, and stores hot pages in BF16 and cold pages as packed INT4 values plus scale metadata. The frozen decode path makes at most two FlashInfer calls per layer and merges their log-sum-exp states. Cold pages are expanded into BF16 scratch for attention; that limitation explains why payload savings do not translate one-for-one to peak memory or speed.

PriorityKV page allocation architecture

Evidence, paper, and reproducibility

Resource Purpose
docs/EVIDENCE.md Canonical claim registry and external-audit response
RESULTS.md Frozen metrics and prior result tables
paper/prioritykv_arxiv.tex Handwritten standalone arXiv source
paper/README.md PDF build and packaging instructions
docs/DATASET.md PriorityBench-A task and split specification
docs/REPRODUCIBILITY.md Local, artifact, and H200 reproduction levels
FINAL_RUN_MANIFEST.yaml Frozen model, benchmark, config, and job IDs
docs/PRAJNA_BFCL_HANDOFF.md Slurm/H100 handoff for external BFCL + CPU τ-bench audit
docs/PRAJNA_CURSOR_PROMPT.md Copy-paste prompt for the Cursor agent on Prajna

Repository layout

src/prioritybench/    deterministic benchmark generator and scorers
src/prioritykv/       role policies, packed cache, and FlashInfer decode
configs/              frozen experiment configurations
jobs/                 canonical H200 commands and result bundles
paper/                arXiv source, readable manuscript, and generated figures
scripts/              reproduction, audit, and figure-generation entrypoints
tests/                sub-project-local CPU and GPU contract tests

Local reproduction

Install the CPU development environment and run the project-local checks:

git clone https://github.com/Arush777/Priority_KV.git
cd Priority_KV
./scripts/sync.sh
./scripts/check.sh

Regenerate and audit the locked benchmark:

PYTHONPATH=src uv run python scripts/mk_bench.py --mode w3_lock
PYTHONPATH=src uv run python scripts/audit_bench.py

The camera-ready paper uses the reviewed PNG assets tracked in paper/figures/. The older deterministic SVG/PDF generator remains available for experiment-side reference plots, but its output is not the camera-ready figure set used by the manuscript.

Build the paper PDF:

cd paper
mkdir -p build
tectonic -X compile prioritykv_arxiv.tex --outdir build --keep-logs

GPU reproduction

GPU dependencies are isolated from the CPU environment:

./scripts/sync.sh --cuda
export PRIORITYKV_SCRATCH=/data/anupam/scratch/prioritykv

Canonical commands and device assignments are indexed in FINAL_RUN_MANIFEST.yaml. Do not run GPU code on a login node, and use no more than two H200 GPUs per job.

Limitations

  • PriorityBench-A is synthetic and agent-specific; it is not LongBench or RULER.
  • The tagger uses visible role and template-like markers. Buried-state controls expose the boundary: structure loses to FullKV on all three slices.
  • Attention baselines are matched-budget repository reimplementations, not reference-code runs; chunked H2O has a fidelity caveat.
  • The Qwen structure–SnapKV difference is not significant (McNemar p=0.125).
  • Llama at 5% keep reverses the Qwen ordering on both evaluated slices.
  • Systems measurements are single-request; they do not establish serving throughput, concurrency behavior, or tail latency.
  • Streamed-cold is a smoke test, and BF16 cold scratch limits peak-memory savings.

Citation and license

Citation metadata is in CITATION.cff. PriorityKV is licensed under the Apache License 2.0. Model weights, benchmark dependencies, and third-party libraries retain their respective licenses. Author affiliations do not imply institutional endorsement.

Contributing

Read CONTRIBUTING.md before changing benchmark semantics, frozen claims, or canonical run configurations. Security reports should follow SECURITY.md.

About

PriorityKV: structure-aware KV-cache eviction for long agent traces, with matched-budget Qwen3-8B evaluation and separate BF16/INT4 memory–latency measurements on H200.

Resources

Contributing

Security policy

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Used by

Contributors

Languages