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.
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 | |
|---|---|---|---|
| Max concurrent requests @ 8192 context, zero errors/timeouts | 1 | 4 | 4x |
| Max concurrent requests @ 16384 context | 0 (refuses to boot) | 2 | boots+serves where plain can't even start |
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.
- 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 latency — at batch 1, STAR-KV's decode kernel is slower than FlashAttention-2 (launch-overhead bound, 0.36-0.87x); it crosses over to faster than FA2 around batch 8-16 and stays 1.2-1.46x faster at real serving batch sizes.
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
python -m proxima_vllm.calibration.distill \
--base-model lmsys/longchat-7b-v1.5-32k \
--output ./checkpoints/longchat-7b-star-kv \
--profiles aggressive,balanced,conservative
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.
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. This section keeps the headline wins; full raw sweep logs (including the messier kernel-level tradeoffs) live in git history if you need them.
Same GPU, same model, distinct (non-shared-prefix) prompts per concurrent request,
enable_prefix_caching=False
, fixed 512-token output, 120s request timeout. Capacity = serves
every request without erroring or timing out. max_concurrency
= highest concurrency tested with zero errors before the next doubling fails.
| context | plain vLLM max concurrency | STAR-KV max concurrency | delta |
|---|---|---|---|
| 8192 | 1 | 4 | |
| 4x | |||
| 16384 | 0 (boot refused) | 2 | |
| STAR-KV serves here; plain can't boot at all |
Detail per level:
8192: plain vLLM — conc=1 ok (13.1 tok/s, ttft 2.5s); conc=2 both requests time out. Boot log:Available KV cache memory: 6.99 GiB
, 14,304 tokens. STAR-KV — conc=1/2/4 all ok (up to 20.2 tok/s, ttft 9.7s at conc=4); conc=8 all time out. Boot log:Available KV cache memory: 7.05 GiB
, 44,800 tokens.16384: plain vLLM refuses to boot at everygpu_memory_utilization
tried (0.9/0.85/0.8/0.75):ValueError: ... 8.26 GiB KV cache is needed, which is larger than the available KV cache memory (6.99 GiB)
. STAR-KV boots (gpu_memory_utilization=0.85
, 5.95 GiB / 37,792 tokens) and serves conc=1/2 ok (up to 9.8 tok/s, ttft 14.1s); conc=4 all time out.32768: plain vLLM boot-refused at every util level (known from the KV-cache-capacity numbers below); STAR-KV concurrency ceiling at this context not yet measured under this protocol.
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 |
At 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 — 20 concurrent requests at 178.6 tok/s.
Across the full context sweep (prompts sized to max_model_len - 256
):
| engine | 4096 | 8192 | 16384 | 32768 |
|---|---|---|---|---|
| plain vLLM | ok | ok | boot refused | |
| boot refused | ||||
| STAR-KV | ok | ok | ok | ok |
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.
scripts/bench/bench_attn_step.py
; 32-head MHA, head_size 128, rank 64, fp16. Speedup > 1 means STAR-KV is faster:
| seq len | batch 1 | batch 16 | batch 64 |
|---|---|---|---|
| 1024 | 0.36x | 1.23x | 1.40x |
| 2048 | 0.57x | 1.33x | 1.43x |
| 4096 | 0.84x | 1.40x | 1.43x |
| 8192 | 0.87x | 1.44x | 1.46x |
1.2-1.46x faster than FA2 at the batch sizes real continuous-batching serving actually runs at (>=16); the crossover point is around batch 8-16 regardless of sequence length.
max_model_len=4096
, gpu_memory_utilization=0.85
, sweeping max_concurrency
:
STAR-KV keeps admitting more concurrent decode sequences — mean batch size climbs to ~100-110 before saturating KV cache, vs plain vLLM saturating around batch ~35-47 — and output tok/s converges to near parity (460-483 tok/s) once both are fully saturated, so the capacity win isn't paid for in aggregate throughput at the top end.
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 | 8192 | 7.6 | 525.8 | 1.00 |
| STAR-KV | 9216 | 8192 | 27.8 | 143.9 | 0.98 |
| STAR-KV | 17408 | 16384 | 12.3 | 143.8 | 0.93 |
Plain vLLM saturates KV and preempts heavily under sustained long generation at this concurrency (tpot p99 525.8ms); STAR-KV holds steady tpot p99 (~115-144ms across all tested lengths) including at 17408 context, which plain cannot boot at all.
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.
Apache-2.0. Not affiliated with the STAR-KV paper authors or the vLLM project.