Code for AbridgeKV: Compressing Long Reasoning into the KV Cache of Learned Pause Tokens (Ziyang Jack Cai, Lijie Yang, Vasilis Kontonis, Dimitris Papailiopoulos, Shivam Garg).

AbridgeKV trains a model to compress its own output. After every block of C = 512 generated tokens, the model reads K learned pause tokens: extra input tokens with trained embeddings that produce no output. Once the pause tokens have encoded the block, the block's KV entries are freed and later tokens attend to the pause tokens' entries instead. Three choices make this work:

Qwen3-4B-Base, fine-tuned on 2B response tokens of nvidia/OpenMathReasoning; every arm uses the same data and recipe. Suite: 446 problems from AIME, HMMT, AMC, BRUMO, CMIMC and SMT, 8 samples per problem, 32K-token cap.

Served in vLLM, the per-layer model's cache is 3.5× smaller per sequence, and a saturated H100 finishes a problem 3.2× faster than with full KV.

Path Contents

abridgekv/ model wrapper and training: compress_forward.py (pause tokens, masks, in-block positions), sft_ddp.py (trainer), data_stream.py (fixed data order), gen_compress.py (reference generation), lightthinker.py (baseline)

launch/ one script per step of the pipeline below; shared settings in launch/env.sh

scripts/ data preparation, schedule measurement, export, evaluation, benchmarks, parity checks and tests

third_party/vllm_overlay/ vLLM 0.13.0 overlay that serves the models (README)

results/schedules/ the measured per-layer schedule used in the paper (math, and the coding-domain re-measurement)

paper/ the evaluation records behind the paper's tables and figures, and the scripts that draw them

data/ evaluation suite definition (README); everything else is downloaded

Setup

Two environments: training needs recent transformers; serving needs vLLM 0.13.0, which pins older versions.

training, probes, export

pip install -r requirements.txt

serving and evaluation (separate environment)

pip install vllm==0.13.0 math-verify==0.9.0 datasets bash third_party/vllm_overlay/install_overlay.sh

Outputs go to $OUTPUT_DIR (default ./outputs). The paper's runs used one node of 8× H100-80GB per training run (about 15 hours each) and one H100 or A100-80GB per evaluation shard.

bash launch/prep_data.sh # training pool + evaluation sets

ARM=uniform bash launch/train.sh # 128 code rows per block in every layer RUN_DIR=outputs/runs/uniform_c512_k128 bash launch/measure_schedule.sh # -> perlayer_schedule.json ARM=perlayer bash launch/train.sh # AbridgeKV; SCHEDULE=<json> to use your own schedule ARM=vanilla bash launch/train.sh # full-KV reference ARM=appended bash launch/train.sh # ablation: code rows at appended positions bash launch/train_lightthinker.sh # LightThinker baseline

RUN_DIR=outputs/runs/perlayer_c512_k128 ARM=perlayer bash launch/export.sh # -> $RUN_DIR/hf MODEL=outputs/runs/perlayer_c512_k128/hf ARM=perlayer TAG=perlayer bash launch/eval_suite.sh python3 scripts/summarize_suite.py --root outputs/eval --out outputs/suite.json --md outputs/suite.md

launch/train.sh resumes from the newest complete checkpoint when re-run. The data order is fixed and independent of the GPU count, so runs on other hardware see the same examples at every step. The measured schedule is checked in (results/schedules/sched_trunc_perlayer_mean128.json), so the per-layer arm can be trained without the uniform run first.

R-KV. Install the authors' vLLM port (Zefan-Cai/R-KV) in its own environment and run ARM=rkv RKV_BUDGET=4096 MODEL=<full-KV export> bash launch/eval_suite.sh there. The script sizes the port's scoring buffer from the budget; with a smaller buffer the port silently keeps the full cache.

Speed. bash third_party/install_shadow.sh async $VLLM_SHADOW builds a patched copy of vLLM with asynchronous scheduling, then MODEL=<export> ARM=perlayer TAG=perlayer VLLM_SHADOW=$VLLM_SHADOW bash launch/bench.sh replays 512 requests at the arm's suite generation lengths and records the time to finish.

Tables and figures. cd paper && python3 build_tables.py && python3 plots.py regenerates every table and figure from the records in paper/results.json and paper/data/.

python3 scripts/test_compresslm_vllm_state.py # vLLM scheduler state, CPU python3 scripts/test_compresslm_layer_quotas.py # per-layer quota masks and spans, CPU python3 scripts/test_compresslm_physical_quotas.py # one KV group per quota, CPU python3 scripts/test_sft_head_code_quotas.py # trainer quota handling, CPU python3 scripts/test_math_grading.py # answer grading LIGHTTHINKER_REPO=<clone of zjunlp/LightThinker> python3 scripts/test_lightthinker_mask.py LT_SHADOW=<dir from install_shadow.sh lt> python3 scripts/test_lightthinker_vllm_state.py

On a GPU, scripts/check_compresslm_vllm_parity.py and scripts/check_teacher_forced_parity.py compare vLLM's token log-probabilities with the training-time forward pass, and scripts/check_full_attention_control.py checks that a served model attending to the evicted text would be detected.

@article{cai2026abridgekv, title = {AbridgeKV: Compressing Long Reasoning into the KV Cache of Learned Pause Tokens}, author = {Cai, Ziyang Jack and Yang, Lijie and Kontonis, Vasilis and Papailiopoulos, Dimitris and Garg, Shivam}, year = {2026} }

License

Apache-2.0. The vLLM overlay derives from vLLM (Apache-2.0) and the microsoft/memento overlay (MIT); see NOTICE.