⚡ Proxima

Low-rank KV Cache Compression for vLLM

An out-of-tree vLLM plugin implementing STAR-KV — low-rank KV cache compression via learned soft-thresholding for adaptive rank control Bhatnagar et al., ICML 2026 Spotlight .

STAR-KV replaces a transformer's full-dimension key/value cache with a per-layer low-rank projection (head-wise decomposition for K, joint decomposition for V), calibrated offline via short knowledge distillation, and stored in a compressed paged format. This repo reproduces that pipeline and wires it into vLLM's V1 continuous-batching, paged-attention serving stack via dedicated Triton kernels.

Why

KV-cache compression: fit more concurrent requests, and longer contexts, in the same GPU memory. Measured end to end on real GPU hardware (Modal L4 24GB):

	plain vLLM	STAR-KV	delta
GPU KV-cache blocks, same memory budget (max_model_len=4096)	756	2353	3.11x
Max max_model_len this L4 can boot at	8192	32768	boots where plain refuses

Full breakdown, methodology, and every number behind this table is in Measured results below.

When to use it

• You are memory-bound, not latency-bound: you need more concurrent sequences or longer contexts than stock vLLM fits on the GPU you have, and can tolerate slower per-token decode in exchange for not OOMing or refusing to boot.
• You are running long-context workloads (16k-32k+ tokens) on GPUs where plain vLLM's boot-time KV-cache pre-flight check refuses to start.
• You are not chasing single-request or low-concurrency latency — at batch 1, STAR-KV's decode kernel is slower than FlashAttention-2 (launch-overhead bound), and end-to-end throughput at moderate concurrency currently trails plain vLLM until the batch sizes and contexts get large enough for the capacity win to dominate.

Install

Not on PyPI. Install from source:

git clone https://github.com/Tenosra/Proxima
cd Proxima
pip install -e .   # pins vllm==0.10.1.1, see docs/vllm_patch_notes.md for why

Quick start

# 1. Calibrate a STAR-KV checkpoint from a base HF model (offline, one-time, no vLLM involved)
python -m proxima_vllm.calibration.distill \
    --base-model lmsys/longchat-7b-v1.5-32k \
    --output ./checkpoints/longchat-7b-star-kv \
    --profiles aggressive,balanced,conservative

# 2. Serve it with vLLM
vllm serve ./checkpoints/longchat-7b-star-kv \
    --additional-config '{"proxima_vllm": {"kv_compression": "star_kv", "rank_profile": "balanced", "use_triton_kernels": true}}'

The reference calibrated checkpoint used for the numbers below is published (private) at iampoppyxx/longchat-7b-starkv-proxima on the HF Hub.

Measured results

Baseline is vLLM 0.10.1.1's own V1 FlashAttentionBackend. On L4 (sm89/Ada) vLLM falls back to FlashAttention-2, not FA3 (FA3 requires Hopper) — that's what every number below is compared against.

KV-cache capacity

balanced profile, same model/checkpoint, same memory budget (max_model_len=4096, gpu_memory_utilization=0.85, block size 16):

	GPU blocks	vs stock vLLM
plain vLLM	756	1x
STAR-KV balanced	2353	3.11x

Long-context capacity

max_model_len=32768: plain vLLM's boot-time pre-flight check refuses to start at any gpu_memory_utilization (needs 16.00 GiB KV cache, only ~5.45 GiB available). STAR-KV boots at gpu_memory_utilization=0.83 and serves real traffic:

requests	context	output tok/s	wall (s)	ttft p50 (ms)
20	32768	178.6	27.4	2578.3

Per-decode-step attention

scripts/bench/bench_attn_step.py; 32-head MHA, head_size 128, rank 64, fp16. Median ms, speedup > 1 means STAR-KV is faster:

seq len	batch	vLLM FA2 (ms)	STAR-KV (ms)	speedup
1024	1	0.150	0.414	0.36x
1024	16	1.254	1.022	1.23x
1024	64	4.465	3.190	1.40x
2048	1	0.247	0.430	0.57x
2048	16	2.314	1.735	1.33x
2048	64	8.718	6.079	1.43x
4096	1	0.465	0.556	0.84x
4096	16	4.438	3.160	1.40x
4096	64	17.279	12.092	1.43x
8192	1	0.736	0.849	0.87x
8192	16	8.708	6.064	1.44x
8192	64	34.223	23.520	1.46x
32768	20	42.636	29.968	1.42x

Faster at batch >= 16 (1.2-1.46x), slower at batch 1 (0.36-0.87x) — STAR-KV needs 4 kernel launches plus a matmul per step versus FlashAttention's 2, so single-request decode is launch-overhead bound rather than compute bound.

Per-prefill-chunk attention

Same script with --num-new N, one chunked-prefill chunk of N tokens per request. prefill_attn_triton (dedicated tiled kernel) vs decode_attn_triton (shared general path) vs vLLM FA2:

seq len	batch	chunk	vLLM FA2 (ms)	STAR-KV prefill (ms)	STAR-KV decode kernel (ms)	prefill vs decode kernel	vs FA2
8192	1	128	0.898	1.432	1.800	1.26x	0.63x
8192	4	128	2.662	3.802	4.911	1.29x	0.70x
32768	1	128	3.180	3.797	3.810	1.00x	0.84x
32768	4	128	9.845	11.498	14.112	1.23x	0.86x
8192	1	512	1.856	3.641	4.733	1.30x	0.51x
8192	4	512	7.945	12.831	16.418	1.28x	0.62x
32768	1	512	6.714	10.853	13.385	1.23x	0.62x
32768	4	512	27.148	36.975	50.930	1.38x	0.73x

Dedicated prefill kernel is 1.2-1.4x faster than routing through the decode kernel, but still 0.51-0.86x of FA2 — prefill reconstructs K from the low-rank basis per query tile, FlashAttention reads dense K/V it never had to compress.

End-to-end serving

scripts/sharegpt_bench/serving_benchmark.py; 20 concurrent ShareGPT requests, max_tokens=256, rank_profile: balanced:

context	engine	output tok/s	wall (s)	ttft p50 (ms)
4096	plain vLLM	292.5	17.5	189.1
4096	STAR-KV balanced	177.3	27.6	2643.1

STAR-KV is 1.65x slower than plain vLLM at max_model_len=4096 despite winning on both the decode kernel (1.44x at this batch size) and KV-cache capacity (3.11x) — per-step kernel and per-chunk prefill wins do not net out to an end-to-end win here; reconstruction cost, TTFT, and scheduling overhead still dominate at this concurrency and context length.

Kernel profiling breakdown

torch.profiler on the same 20-request, max_tokens=256, max_model_len=4096 workload, self-CUDA time compared between engines:

kernel	plain vLLM	STAR-KV	avg per call
attention forward	flash_fwd_splitkv_kernel, 8160 calls, 8.85% (1.384s)	_starkv_attn_kernel, 8160 calls, 26.84% (6.587s)	169.6us vs 807.2us
cache write	reshape_and_cache_flash_kernel, 8192 calls, 0.26% (39.9ms)	_write_kv_kernel, 8192 calls, 16.14% (3.962s)	4.9us vs 483.6us
combine (split-K only)	n/a	_combine_splits_kernel, 8160 calls, 0.16% (38.9ms)	— vs 4.8us

The gap is concentrated in real per-call kernel cost at the batch sizes async serving actually produces (not Python glue, host syncs, or un-fused reconstruction — those are already eliminated). _write_kv_kernel does real compression compute (low-rank projection + normalization) where reshape_and_cache_flash is a bandwidth-bound copy — a structural cost of compression, not a fixable launch inefficiency.

Campaign: max sustainable concurrency

max_model_len=4096, gpu_memory_utilization=0.85, sweeping max_concurrency. batch is mean decode batch size; kv used is peak KV-cache occupancy fraction (1.00 = saturated):

engine	conc	out tok/s	ttft p50 (ms)	tpot p99 (ms)	batch	kv used
plain	1	18.4	66	57.6	1.0	0.09
plain	4	61.7	125	59.5	3.5	0.15
plain	8	113.5	170	73.2	6.8	0.23
plain	16	205.1	291	67.1	12.9	0.47
plain	32	342.9	301	83.9	25.1	0.80
plain	64	429.9	1044	297.0	35.2	1.00
plain	128	478.1	11764	709.4	47.4	1.00
plain	256	483.5	50722	825.6	47.9	1.00
STAR-KV	1	6.9	119	407.9	1.0	0.03
STAR-KV	4	9.8	224	1394.7	3.4	0.06
STAR-KV	8	61.2	235	227.4	5.8	0.09
STAR-KV	16	121.8	302	172.9	12.7	0.18
STAR-KV	32	210.0	523	275.6	25.4	0.32
STAR-KV	64	325.5	1063	500.4	50.9	0.64
STAR-KV	128	469.2	2014	540.8	99.9	1.00
STAR-KV	256	460.9	4566	715.2	109.5	1.00

Plain vLLM saturates around concurrency 64 (batch ~35); STAR-KV keeps admitting more concurrent decode sequences up to ~2.3x the batch size (99.9-109.5 vs 47.4-47.9) before saturating, and output tok/s converges to near parity once both are saturated.

Campaign: max context length per engine

Prompts sized to nearly fill each context window (target_prompt_tokens = max_model_len - 256), max_num_batched_tokens=2048:

engine	context	status	out tok/s	ttft p50 (ms)
plain	4096	ok	22.1	3492
plain	8192	ok	3.2	107049
plain	16384	boot refused	—	—
plain	32768	boot refused	—	—
STAR-KV	4096	ok	24.6	2388
STAR-KV	8192	ok	2.1	164620
STAR-KV	16384	ok	0.8	159611
STAR-KV	32768	ok	0.6	447585

Plain vLLM refuses to boot above 8192 on this L4; STAR-KV boots and serves at every context up to the model's 32768 limit.

Campaign: decode-dominated long-generation

Single long-generation request(s), concurrency=4 at max_model_len=9216, concurrency=2 at 17408:

engine	context	max_tokens	out tok/s	tpot p99 (ms)	kv used
plain	9216	1024	57.8	68.8	0.68
plain	9216	2048	47.9	82.8	1.00
plain	9216	4096	20.2	197.6	1.00
plain	9216	8192	7.6	525.8	1.00
STAR-KV	9216	1024	5.7	395.2	0.22
STAR-KV	9216	2048	34.5	115.8	0.33
STAR-KV	9216	4096	32.0	124.7	0.54
STAR-KV	9216	8192	27.8	143.9	0.98
STAR-KV	17408	16384	12.3	143.8	0.93

Plain vLLM saturates KV (1.00) and preempts heavily past max_tokens=2048 at this concurrency; STAR-KV stays under 1.00 throughout and holds steady tpot_p99 (~115-144ms) across all lengths including 17408, a context plain cannot boot at all.

Campaign: scheduler efficiency

Mean decode batch size and fraction of decode steps stuck at batch=1, aggregated across all phases:

engine	mean decode batch	steps at batch=1
plain	5.80	32.5%
STAR-KV	7.20	12.1%

STAR-KV spends less time decode-batch-starved (batch=1) across the campaign — consistent with its larger effective KV capacity keeping more requests concurrently resident.

Scope

Verified end to end on lmsys/longchat-7b-v1.5-32k, real ShareGPT serving traffic, on Modal L4 GPU hardware. Multi-model support (Qwen/Mistral/Gemma/Llama-3) is written but not yet validated on hardware beyond this one model.

License

Apache-2.0. Not affiliated with the STAR-KV paper authors or the vLLM project.