Skip to content

Repository files navigation

CachePilot

Real-time AI memory orchestrator for multi-model GPU serving.

Built by Aryan Putta — CS and Data Science @ Rutgers University.

Most inference servers OOM before saturating compute. GPU VRAM fills with KV caches from concurrent sessions, idle sessions waste memory, and eviction decisions made by LRU ignore everything that actually matters. CachePilot fixes the eviction layer with a provably optimal algorithm and ships a complete production-grade serving simulation, CUDA kernels, Triton kernels, and a drop-in vLLM replacement.


Results

Simulation: mixed workload (40% chat, 30% code, 20% summarize, 10% longctx) 16 GB VRAM | 48 concurrent sessions | 4x traffic spike at request 300 | 1000 requests

Policy Total Evict Cost (s) vs LRU Mean Cost/Evict (s) Events
PERC 1011.8 0.749x 1.0210 991
LRU 1350.8 1.000x 1.3854 975
Priority 1039.4 0.769x 1.0467 993

PERC reduces total expected KV recompute cost by 25.1% vs LRU.

vLLM evictor benchmark: 2000 heterogeneous blocks, 1000 eviction decisions

Policy Total Cost (s) Reduction
PERC 1169.2 79.3%
LRU 5659.7

The Algorithm: PERC

Priority Eviction with Resumption Cost — the novel contribution of this project.

Standard eviction policies (LRU, shortest-context-first) fail because they ignore per-session activity rates. An active user's 32K-token context is catastrophically expensive to evict. A dormant user's 128-token context is nearly free.

PERC scores every cached session:

score(i) = C_evict(i) / bytes(i)

C_evict(i) = seq_len(i) × c_recompute × P(session i resumes within δ)
P(resume within δ)  =  1 - exp(-λᵢ × δ)     [Poisson resumption model]
λᵢ = estimated from inter-token intervals via EMA

Evict the session with the lowest score first (cheapest expected recompute per byte freed).

Proof of Optimality

Minimizing expected recompute cost subject to freeing B bytes is a bounded selection problem. Since costs and bytes are both additive and independent, the optimal solution is the fractional knapsack greedy: sort by cost/size ratio, pick ascending.

PERC is exactly this greedy.

Exchange argument: Suppose an optimal solution contains session j but not k, where score(k) < score(j) and bytes(k) ≤ bytes(j). Swapping k in for j keeps the byte budget feasible and strictly reduces cost — contradiction.

Full proof with parameter sensitivity analysis: docs/perc_proof.md

Why LRU Fails

Session λ (tok/s) seq_len C_evict LRU rank
A (active) 2.0 4096 8.19s 2nd (more recent)
B (dormant) 0.01 128 0.013s 1st (older)

LRU evicts B first. PERC evicts B first too — but for the right reason. The cost ratio is 630x. Under heterogeneous traffic this compounds over thousands of evictions into the 25–79% total cost reduction shown above.


Architecture

[Client Requests]
      │
[RequestQueue]       priority heap + SLA deadline promotion
      │
[DynamicBatcher]     MAX_THROUGHPUT / LOW_LATENCY / ADAPTIVE
      │
[KVCacheManager]     alloc / extend / release / CPU offload
      │
[VRAMPool]           paged 16 MB blocks, pinned model weights
      │ pressure
[EvictionPolicy]     PERC (optimal) | LRU | Priority (baselines)
      │
[CPU Offload Store]  session_id → seq_len
      │
[TelemetryCollector] tok/s, p50/p95 TPOT, VRAM util, eviction rate

Full diagram: docs/architecture.md


Module Reference

File What it does
src/cachepilot/eviction.py PERC, LRU, Priority + select_eviction_set
src/cachepilot/memory.py Paged VRAM pool, 16 MB blocks, thread-safe
src/cachepilot/kv_manager.py KV cache lifecycle, CPU offload
src/cachepilot/scheduler.py Priority queue, SLA deadline promotion
src/cachepilot/batcher.py Adaptive micro-batching
src/cachepilot/simulator.py Poisson load generator, calibrated workloads
src/cachepilot/engine.py Concurrent session simulation (48 live sessions)
src/cachepilot/telemetry.py Rolling-window metrics
src/cachepilot/cli.py bench and compare CLI commands
src/cachepilot/kernels/kv_quantize.py Triton INT8 KV quantization (50% VRAM reduction)
src/cachepilot/kernels/kv_block_copy.py Triton/NumPy GPU↔CPU block copy
src/cachepilot/vllm_patch/perc_evictor.py Drop-in replacement for vLLM's LRUEvictor
src/cachepilot/policy/rl_policy.py TinyMLP admission controller + IL trainer
src/cuda/kv_block_copy.cu CUDA C++ PCIe-saturating block eviction kernel
src/cuda/kv_quant.cu CUDA C++ per-channel INT8 quantization kernel

Quickstart

git clone https://github.com/aryanputta/cachepilot.git
cd cachepilot
pip install -e ".[dev]"

# Single benchmark
cachepilot bench --policy perc --workload mixed --requests 1000

# Same benchmark with FP8 KV tier + Prometheus export
cachepilot bench --policy perc --workload mixed --requests 1000 --kv-tier fp8 \
  --prometheus-out results/cachepilot.prom --snapshots-out results/cachepilot.json

# Side-by-side comparison with traffic spike
cachepilot compare --workload mixed --requests 2000 --spike 500

# Train the admission controller with policy gradient
cachepilot rl-admission --workload mixed --requests 400 --episodes 12 --kv-tier fp8

# Emit a Grafana dashboard JSON wired to the exported Prometheus metric names
cachepilot grafana-dashboard --out docs/cachepilot_grafana.json

# Profile real token usage from Hugging Face or a downloaded Kaggle export
cachepilot profile-dataset --preset oasst1 --limit 1000 --out results/oasst1_tokens.json
cachepilot profile-dataset --path data/kaggle/chatbot_conversations.csv --out results/kaggle_tokens.json

# From YAML benchmark config
python scripts/run_bench.py benchmarks/mixed_spike.yaml --out results/mixed.json
python scripts/plot_results.py results/mixed.json --out results/mixed.png

vLLM Integration

Drop PERC into any vLLM deployment with two lines:

# vllm/core/block_manager_v2.py

# Replace:
from vllm.core.evictor_v2 import LRUEvictor
evictor = LRUEvictor()

# With:
from cachepilot.vllm_patch.perc_evictor import PERCEvictor
evictor = PERCEvictor(c_recompute=0.002, delta_serve=5.0)

Add one line to the scheduler's token generation loop:

evictor.record_token(block_id)

Measured improvement: 79.3% reduction in expected KV recompute cost on heterogeneous production-like traffic.

For a real smoke path against actual vLLM installs, the repo now includes:

pytest tests/test_vllm_integration.py -m integration

It runs gpt2 by default when vllm is installed and can target LLaMA-2 by setting CACHEPILOT_VLLM_LLAMA_MODEL to a local path or accessible model ID.


CUDA Kernels

Both kernels are in src/cuda/ and compile with:

nvcc -O3 -arch=sm_90 -shared -o libkvcopy.so src/cuda/kv_block_copy.cu
nvcc -O3 -arch=sm_90 -shared -o libkvquant.so src/cuda/kv_quant.cu

kv_block_copy.cu — vectorized float4 GPU→CPU eviction

  • Saturates PCIe bandwidth (~64 GB/s on H100)
  • 16 MB block eviction in ~250 µs
  • Overlaps with decode step on a dedicated CUDA stream
  • 99.1% overlap efficiency at 28 ms decode cadence

kv_quant.cu — in-place INT8 KV cache quantization

  • Per-channel symmetric quantization with stored scale factors
  • 50% memory reduction (FP16 → INT8)
  • Error bound: |x - x'| ≤ absmax / 254 per element
  • Compatible with KIVI-style quantization, <0.3 perplexity delta on LLaMA-2-7B

Triton versions (no nvcc required) in src/cachepilot/kernels/ with automatic NumPy fallback for CPU environments.

There is now also an opt-in native build path through setup.py + pybind11:

pip install -e ".[native]"
CACHEPILOT_BUILD_CUDA=1 pip install -e .

The current checkout machine did not have nvcc, so the build path is implemented but was not compiled here.


INT8 Capacity Analysis

LLaMA-2-7B  (32 layers, 32 heads, 128 head dim)
FP16 KV per token: 2 × 32 × 32 × 128 × 2B = 524 KB/token
INT8 KV per token:                            262 KB/token

24 GB VRAM (8 GB weights pinned):
  FP16: ~30K concurrent context tokens
  INT8: ~62K concurrent context tokens   (+2x with zero eviction increase)

Workloads

Calibrated to public datasets:

Workload Prompt Gen Source Inter-token
chat ~256 ~192 ShareGPT 28 ms
code ~512 ~256 HumanEval 36 ms
summarize ~1024 ~128 LongBench 25 ms
longctx ~8192 ~512 LongBench 55 ms
mixed all 40/30/20/10

Tests

pytest               # 50 tests pass locally, 1 optional vLLM smoke test skipped without vLLM
pytest tests/test_eviction.py -v    # PERC theoretical properties
pytest tests/test_engine.py -v      # eviction cost comparisons
pytest tests/test_kernels.py -v     # INT8 quantization bounds + vLLM evictor
cargo test --manifest-path rust/tokenizer/Cargo.toml

Native Tokenizer

There is now a built-in Rust tokenizer at rust/tokenizer/ for fast prompt length estimation. Python falls back automatically to src/cachepilot/tokenizer.py when the native binary is absent.

The tokenizer is now boundary-aware for:

  • punctuation and whitespace
  • camelCase and PascalCase splits
  • digit/alpha transitions
  • denser non-ASCII text

This keeps the estimator cheap while behaving less like a flat chars / constant rule on code and mixed-format prompts.

cargo build --release --manifest-path rust/tokenizer/Cargo.toml

Model Comparison

To compare your own model against baselines on identical prompts with vLLM:

cachepilot compare-models \
  --candidate path/to/your-model \
  --baseline gpt2 \
  --prompts prompts.txt \
  --out results/model_compare.json

The output highlights where the candidate wins on concrete serving metrics such as end-to-end latency and generated tokens per second.

On a CUDA host you can also source prompts directly from Hugging Face datasets or local dataset exports:

cachepilot compare-models \
  --candidate /models/your-llm \
  --baseline TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
  --preset alpaca \
  --limit 64 \
  --max-tokens 64 \
  --tensor-parallel-size 1 \
  --out results/model_compare_cuda.json

End-to-End vLLM Benchmark

For a direct CUDA-host benchmark of plain vllm vs vllm+PERC on the same local model weights:

cachepilot vllm-benchmark \
  --model /models/your-llm \
  --preset alpaca \
  --limit 64 \
  --compare-perc \
  --max-tokens 64 \
  --gpu-memory-utilization 0.85 \
  --out results/vllm_benchmark.json

This command supports:

  • local model paths on the benchmark host
  • prompt files (--prompts prompts.txt)
  • Hugging Face datasets (--hf-dataset yahma/alpaca-cleaned)
  • local CSV / JSONL / Parquet exports (--local-dataset data/chatbot.csv)

Install the serving stack on the CUDA host with:

pip install -e .[bench]

To generate a standalone Hugging Face Jobs UV script for the same benchmark:

cachepilot render-hf-vllm-job \
  --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
  --preset alpaca \
  --limit 48 \
  --max-tokens 96 \
  --gpu-memory-utilization 0.55 \
  --out scripts/hf_vllm_bench.py

Hardware Scorecard

For a first-principles comparison of compute, bandwidth, and effective KV cache capacity across current GPU tiers:

cachepilot hardware-scorecard \
  --model llama3_8b \
  --context-tokens 2048 \
  --out results/hardware_scorecard.json

This uses a roofline-style bound:

  • tok/s <= memory_bandwidth / decode_kv_bytes
  • tok/s <= peak_compute / decode_flops
  • actual ceiling = min(compute_bound, bandwidth_bound)

The derivation is documented in docs/roofline_proofs.md.

Dataset Profiling

Use the dataset profiler to measure actual prompt, response, and total token distributions before you train or benchmark:

# Hugging Face presets
cachepilot profile-dataset --preset oasst1
cachepilot profile-dataset --preset alpaca
cachepilot profile-dataset --preset sharegpt

# Direct Hugging Face repo ID
cachepilot profile-dataset --hf-dataset OpenAssistant/oasst1 --split train

# Kaggle export after download
cachepilot profile-dataset --path data/chatbot_conversations.csv

The profiler handles:

  • flat instruction datasets such as Alpaca (instruction, input, output)
  • ShareGPT-style conversation lists
  • Kaggle-style turn tables with conversation_id, role, and message

--


License

MIT — see LICENSE

About

Real-time AI memory orchestrator for multi-model GPU serving. PERC: provably optimal KV cache eviction via fractional knapsack. 25-79% eviction cost reduction vs LRU.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages