KVPress Inference

August 3, 2026 ยท View on GitHub

Anonymous inference-only utilities for teacher generation and KVPress student-cache compression.

The student path only keeps these KVPress methods:

  • snapkv
  • streamingllm
  • tova
  • knorm
  • lagkv
  • chunkkv
  • adakv
  • pyramidkv

No training code, checkpoints, logs, result tables, local cluster paths, or private dataset paths are included.

Install

pip install -r requirements.txt

KVPress is provided by NVIDIA and can be installed with pip install kvpress.

Teacher Generation

Generate full teacher traces from a JSONL file of problems:

python -m kvpress_inference.teacher_only \
  --data /path/to/problems.jsonl \
  --model_name Qwen/Qwen3-8B \
  --out_jsonl outputs/teacher.jsonl \
  --teacher_decode_tokens 12000 \
  --temperature 0.6 \
  --top_p 0.95 \
  --top_k 20

Each output row contains think, teacher_text, student_text, answer, and meta.

Student KVPress Inference

Run a compressed student answer from a teacher JSONL that already contains think:

python -m kvpress_inference.student_only \
  --data outputs/teacher.jsonl.rank0.jsonl \
  --model_name Qwen/Qwen3-8B \
  --method snapkv \
  --mem_len 256 \
  --out_jsonl outputs/student_snapkv.jsonl \
  --student_decode_tokens 1500

Useful switches:

  • --method: one of snapkv, streamingllm, tova, knorm, lagkv, chunkkv, adakv, pyramidkv
  • --mem_len: reasoning-token KV budget
  • --compress_prompt_kv: let prompt tokens compete for the same KV budget
  • --device_map auto: use Hugging Face/Accelerate model parallel loading
  • --dtype auto: let Transformers choose model dtype

Distributed runs via torchrun are supported. Each rank writes OUT.rank{rank}.jsonl.

Input Format

Teacher generation expects JSONL rows with one of these fields for the problem text:

  • problem
  • question
  • Question

Optional answer fields:

  • answer
  • gt
  • Ground Truth Answer

Student inference expects teacher rows with:

  • think
  • optionally teacher_text

The code has dataset-mode prompt templates for generic math, IFEval, MedCalc, MedXpertQA, RULER, and LiveCodeBench based on the input path/name.