KVLearn

Learning-based KV cache retention for disaggregated LLM serving.

This repository implements the control-plane from the paper To Keep or Not to Keep: Learning KV Cache Retention in Disaggregated LLM Serving Systems:

Component	Role
PRP	Prefix Reuse Predictor — 2-layer MLP (577 params) over 16 structural / temporal / pool-state features
CARS	Cost-Aware Retention Score — P̂·(R−T) − U with hardware-grounded costs
ATC	Adaptive Threshold Controller — PI loop on pool pressure + hit-rate bias

Admission (Alg. 1) and eviction (Alg. 2) run in the KV-pool coordinator path.

Build

Make (no CMake required):

make -j
make test
./build/kvlearn_sim [n_requests] [zipf_alpha] [pool_GiB]

CMake (optional):

cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build -j
ctest --test-dir build --output-on-failure

Example:

./build/kvlearn_sim 3000 1.1 4

Layout

include/kvlearn/   public headers (types, cost model, PRP, CARS, ATC, pool, control plane)
src/               implementations
apps/simulate.cpp  Zipf multimodal trace vs No-Cache / LRU-Pool / KVLearn
tests/             unit checks (Prop. 1, PRP learning, heap eviction, …)

Cost model

|b|   = 2 · ℓ · h_kv · d_h · L · δ
R(b)  ≈ α_bw·L          if L < L×   (bandwidth-bound)
      ≈ α_flop·L²       if L ≥ L×   (compute-bound)
T(b)  = |b| / β
U(b)  = γ · |b| · Δt    γ = R̄·λ / M_S
CARS  = P̂(b)·(R−T) − U
KEEP  ⇔  CARS > θ       (θ from ATC)

Proposition 1: for super-linear prefill (q>1), the optimal reuse threshold P*_H decreases with prefix length — longer visual blocks admit at lower predicted reuse.

Integrating into a serving stack

1. Co-locate kvlearn::KVLearn with the global KV pool coordinator (Mooncake-style P → S → D).
2. On radix hit → lookup(req); on miss after prefill → admit(req, block_id).
3. Feed delayed labels via eviction (y=0/1) and maybe_emit_horizon_labels.
4. Call tick(now) periodically so ATC tracks ρ_pool and PRP trains off-path.

Tensor movement (RDMA / NCCL) stays in your data plane; this library only decides keep vs discard.

Defaults settings

• Features d=16, hidden h=32, η=5e−3, replay |B|=20k, batch every 100 labels
• ATC: ρ*=0.85, Kp=2e−2, Ki=5e−4, h_floor=0.55
• A100 + LLaMA-3-8B calib: α_bw≈0.026 ms/tok, α_flop≈8e−6 ms/tok², L×=512
• Fabric default: InfiniBand HDR 200 Gbps (β=25 GB/s)

License

MIT