kvault

Rust-native KV cache tier for LLM inference engines. Ports LMCache's tiered cache architecture and SGLang's RadixAttention prefix sharing into pure Rust, targeting candle-vllm and mistral.rs.

What it does

• Stores KV cache blocks from finished sequences into host memory (DRAM, NVMe, Redis, S3) so they can be reused across requests
• Matches prefixes via a radix tree over token IDs — when two prompts share a common prefix, the cached KV blocks are reused instead of recomputed
• Eliminates redundant prefill compute: a prefix hit on 5 cached blocks for a 7B model saves ~1,280 tokens × ~15ms/1K ≈ 19ms per request
• Tiered storage with automatic promotion: hot blocks stay in DRAM, warm blocks spill to NVMe, cold blocks ship to Redis/S3
• Distributed coordinator tracks which worker holds which prefix across a multi-node fleet, brokering peer-to-peer KV block transfers
• Python bindings via PyO3 so vLLM/LMCache Python users can drop in kvault_py.KvaultEngine as a cache backend
• Zero-copy wire format (84-byte header + raw bytes) with SHA-256 checksums and optional CacheGen compression

Architecture

┌──────────────────────────────────────────────────┐
│  KvaultEngine  (store / retrieve / lookup)      │
├───────────────┬──────────────────┬───────────────┤
│  kvcore       │  gpu             │  storage      │
│  ──────────── │  ──────────────  │  ──────────── │
│  CacheKey     │  KvConnector     │  DramBackend  │
│  TokenDB      │  CandleVllmConn  │  DiskBackend  │
│  RadixTree    │  MistralRsConn   │  RedisBackend  │
│               │                  │  S3Backend    │
└───────────────┴──────────────────┴───────────────┘
                         │
            ┌────────────┼────────────┐
            ▼            ▼            ▼
       crates/rpc   crates/server  crates/serde
       coordinator  /health        wire format
       lookup RPC   /metrics       CacheGen
                    /v1/cache/...

Quick start

cargo build
cargo test

No GPU required — the mock connector validates the full pipeline without hardware.

Rust API usage

use kvault_engine::{EngineConfig, KvaultEngine};
use kvault_gpu::connector::{KvShape, MockGpuConnector};
use kvault_kvcore::ExtraKeys;
use kvault_storage::cpu::DramBackend;
use std::sync::Arc;

#[tokio::main]
async fn main() {
    // Configure for Qwen2.5-0.5B (24 layers, 14 KV heads, 64 head_dim)
    let shape = KvShape {
        num_kv_heads: 14, key_head_dim: 64, value_head_dim: 64,
        block_size: 256, num_layers: 24, flash_layout: true,
        is_mla: false, kv_lora_rank: None, qk_rope_head_dim: None,
    };

    let storage = Arc::new(DramBackend::new(1024 * 1024 * 1024)); // 1 GB
    let gpu = Arc::new(MockGpuConnector::new(shape, 2));
    let engine = Arc::new(KvaultEngine::new(
        EngineConfig::default(), storage, gpu,
        ExtraKeys { model_id: "Qwen/Qwen2.5-0.5B".into(), .. },
        24,
    ));

    // Store finished sequence
    let tokens: Vec<u32> = (0..512).collect();
    engine.store(&tokens, &[0, 1]).await.unwrap();

    // Check cacheability
    let results = engine.lookup(&(0..1024).collect::<Vec<_>>()).await.unwrap();
    println!("First 2 chunks cacheable: {}/{}", results[0].1, results[1].1);

    // Retrieve prefix
    let result = engine.retrieve(&(0..512).collect::<Vec<_>>()).await.unwrap();
    println!("Matched {} tokens from cache", result.matched_tokens);
}

Framework integration

Framework	Status
candle-vllm	Fork with get_kv_blocks/set_kv_blocks hooks; real GPU round-trip test (cargo test -p kvault-gpu --features candle-vllm)
mistral.rs	MistralRsConnector scaffold behind the same KvConnector trait; needs mistralrs-core dep

Project structure

crates/
├── kvcore/     CacheKey, ChunkedTokenDatabase, RadixTree
├── memobj/     MemoryObj (zero-copy bytes, metadata)
├── storage/    StorageBackend trait + DramBackend, DiskBackend, Redis, S3,
│               StorageManager, RooflineModel
├── gpu/        KvConnector trait + MockGpuConnector, CandleVllmConnector
├── engine/     KvaultEngine — store/retrieve/lookup/get_blocks
├── serde/      KvBlockHeader (84-byte wire), CacheGen compression
├── server/     Axum HTTP + Prometheus /metrics, /health, /v1/cache/...
├── rpc/        Framed TCP lookup RPC, distributed coordinator
├── pyo3/       Python bindings (PyO3)
└── benches/    Benchmarking CLI

Status

Phase	Description	Status
0	GPU round-trip on candle-vllm fork	✅ Done
1	kvcore + memobj	✅ Done
2	Single-node cache path	✅ Done
3	NVMe tier	✅ Done
4	Remote tier (Redis, S3) + roofline	✅ Done
5	Serde + CacheGen	✅ Done
6	Control plane (HTTP + RPC)	✅ Done
7	Distributed coordinator	✅ Done
8	PyO3 Python bindings	✅ Done

Tests: 78 passing, 0 failures.

License

MIT — see LICENSE.

Related resources

• Learning resources — papers, blog posts, and docs on KV caches, PagedAttention, RadixAttention, and LMCache
• Full implementation plan
• Grafana dashboard — Prometheus-based observability