Skip to content

DaoCloud/ContextStore

Repository files navigation

ContextStore

Tiered shared storage for LLM inference KV caches.

ContextStore extends the KV cache from a single node's GPU HBM out to host memory, local NVMe, and remote JBOF / NVMe-oF pools. It gives vLLM / Dynamo inference workers cross-instance, cross-session KV reuse so long-context and multi-replica deployments avoid repeating prefill on cache-warm prefixes.

The repository is split into two independently built parts, plus a thin plugin layer:

Part Path Language Role
ContextStore Connector src/contextstore/ Python vLLM / Dynamo KVConnector plugin. Owns codec, prefix index, scheduler/worker split, tiered cache, and the KVService client.
KVService kv-service/ Rust Standalone distributed block-storage service for JBOF / NVMe-oF. Exposes gRPC + RDMA data paths.
NIXL plugin nixl-plugin/ C++ / Rust Loads ContextStore as NIXL's CONTEXTSTORE backend under the OBJ descriptor contract.

Architecture

┌─────────────────────────────────────────────────────────┐
│                    vLLM / Dynamo                        │
└────────────────────┬────────────────────────────────────┘
                     │ KVConnector API
                     ▼
┌─────────────────────────────────────────────────────────┐
│         ContextStore Connector  (src/contextstore/)     │
│  ┌───────────┐  ┌──────────┐  ┌──────────────────────┐  │
│  │  Codec    │  │ Prefix   │  │ Scheduler / Worker   │  │
│  │ (INT8 quant)│ │ Index    │  │ metadata + transfer  │  │
│  └───────────┘  └──────────┘  └──────────────────────┘  │
│                       ▼                                 │
│               StorageBackend                            │
│    L1 HostMemory  →  L2 Local / Sharded  →  L3 KVService│
└──────────────────────────────────┬──────────────────────┘
                                   │ gRPC / RDMA
                                   ▼
┌─────────────────────────────────────────────────────────┐
│                KVService  (kv-service/)                 │
│  MemoryTier (L1)  →  StorageTier (L2)                   │
│  IO Executor: Tier A (thread pool) / Tier B (io_uring)   │
│  RDMA server: pre-registered slab, zero-memcpy WRITE    │
└─────────────────────────────────────────────────────────┘

Connector and KVService are built and deployed independently. The Connector handles KV semantics and the vLLM/Dynamo integration; KVService handles shared capacity and high-bandwidth data paths. Neither depends on the other at build time.


KVService object model & metadata

KVService exposes an object store (not a byte-addressable KV): each write produces an object identified by ObjectKey{namespace, object_key}, and each object carries server-generated metadata that clients treat as opaque.

Metadata service (Rust side). MetadataService (kv-service/server/src/metadata.rs) persists per-object records in a shared Redis metadata store. Each BlockMeta includes:

Field Meaning
object_handle Opaque read handle generated by the server; clients round-trip it verbatim
object_generation Content generation; monotonically increases on overwrite of the same key
content_etag Lightweight content-identity token (avoids re-hashing large payloads on every GET)
layout_version Physical layout version; bumped by rebalance / migration without changing content
striping Optional StripingInfo{chunk_size, chunk_devices, chunk_paths, total_size, chunk_locations} for objects striped across multiple NVMe / nodes

Descriptor-based reads. The gRPC surface splits read into two RPCs:

  • LookupObject returns an ObjectDescriptor (identity + layout_version + size + stripe count) and a PlacementDescriptor (which node / device / RDMA endpoint holds each stripe). No payload bytes are returned.
  • ReadByDescriptor / ReadByDescriptorStream takes a client-cached descriptor and returns the payload. The server verifies the descriptor first; if the object was rewritten or relocated, the server replies FAILED_PRECONDITION and the client re-runs LookupObject.

Speculative concurrent lookup + read (client side). KVClient.get_stream_chunks_cached (src/contextstore/kvservice_client/client.py) hides the round-trip cost of validation:

  • On a cache miss, do LookupObjectReadByDescriptor sequentially and cache the result.
  • On a cache hit, fire LookupObject and ReadByDescriptor(cached_descriptor) concurrently. If identity matches, return the read; if only layout_version changed but content_etag matches, still return the read and refresh the cache; only on content mismatch does the client re-read using the fresh descriptor.

This keeps hot GETs to a single round trip while remaining correct under concurrent overwrite / rebalance.

Multi-node placement. For striped objects the coordinator uses PutPlacementChunk / ReadPlacementChunk / DeletePlacementChunk to write individual stripes directly to the data node that owns them, keyed off PlacementChunk{stripe_index, node_id, grpc_endpoint, rdma_endpoint, device_id, storage_handle, offset, length}. Clients that hold a fresh PlacementDescriptor can bypass the coordinator and read the primary node directly — this is what unlocks the RDMA fast path.


Quick start

1. Install the Python library

git clone git@github.com:DaoCloud/ContextStore.git
cd ContextStore
pip install -e .            # Connector + KVService client
pip install -e '.[test]'    # adds pytest / fakeredis
# Only needed after editing kv-service/proto/kv_service.proto.
pip install -e '.[proto]'   # adds grpcio-tools for `make proto`

2. Build and run KVService (optional L3 tier)

# Build server / client-rs / rdma-ffi with the deployment feature set.
make build

# Start the server (listens on :50051)
./target/release/contextstore-server \
    --config kv-service/configs/server.toml

KVService reads one TOML config file and requires a reachable Redis metadata store. See kv-service/configs/README.md for the config file format and examples.

make build enables the RDMA data path, Tier B io-uring, and Prometheus metrics. It produces the server, Rust client SDK, and RDMA C ABI under target/release/. Regenerate Python protobuf bindings only after changing the protocol:

make proto

3. Wire the Connector into vLLM

Pass a --kv-transfer-config to vLLM v1:

{
  "kv_connector": "ContextStoreConnector",
  "kv_connector_module_path": "contextstore.connector",
  "kv_role": "kv_both",
  "kv_connector_extra_config": {
    "model_id": "Qwen2.5-32B",
    "kv_service_endpoint": "10.0.0.1:50051",
    "host_memory_capacity_gb": 128,
    "kv_service_parallel_channels": 1,
    "rdma_enabled": true,
    "rdma_server_addr": "10.0.0.1:50053",
    "rdma_device": "mlx5_0"
  }
}

All keys map directly to fields on ContextStoreConfig (src/contextstore/core/config.py). Multi-endpoint HA is enabled by passing kv_service_endpoints: ["host1:50051", "host2:50051"] instead of a single endpoint.

4. Or wire into Dynamo

Dynamo loads the same Connector under the KVBM contract, with a contextstore.* dotted-alias layer:

{
  "kv_connector": "DynamoConnector",
  "kv_connector_module_path": "contextstore.integrations.dynamo.connector",
  "kv_role": "kv_both",
  "kv_connector_extra_config": {
    "contextstore.model_id": "Qwen2.5-32B",
    "contextstore.endpoint": "10.0.0.1:50051",
    "contextstore.host_memory_capacity_gb": 128,
    "contextstore.rdma_enabled": true
  }
}

Both dotted (contextstore.foo) and nested ({"contextstore": {"foo": ...}}) forms are accepted; see _DOTTED_ALIASES in src/contextstore/integrations/dynamo/connector.py for the full list.


Repository layout

ContextStore/
├── src/contextstore/       # Python library (Connector + KVService client SDK)
├── kv-service/             # Rust KV Service (server, client-rs, rdma-ffi, configs, deploy)
├── nixl-plugin/            # NIXL CONTEXTSTORE backend (C++ plugin + Rust FFI)
├── tests/                  # Python pytest suite (no GPU / Redis required)
├── Makefile                 # Root build, test, and deployment entry points
├── pyproject.toml
├── CLAUDE.md               # Development conventions
└── LICENSE

Build and test

Python

pip install -e '.[test]'
pytest tests/ -v

Rust

make build                     # server + client-rs + rdma-ffi, deployment features enabled
make server                    # build only contextstore-server
make client-rs                 # build only the Rust client SDK
make rdma-ffi                  # build only the Python ctypes C ABI library
make test-server               # Rust server tests
make test-integration          # requires a running server
make fmt                       # format Rust code
make lint                      # run Rust clippy checks

make build BUILD_TYPE=debug builds the same artifacts with Cargo's debug profile. make proto, make proto-rust, and make proto-python regenerate protocol bindings; the Python targets require grpcio-tools.

Benchmarks

./target/release/cs-kvservice-bench --help
./target/release/cs-bench --help
./target/release/cs-rdma-bench --help
python kv-service/benchmarks/run_benchmark.py --endpoint localhost:50051

Run the Python benchmark wrapper through the root Makefile with:

make bench

Deployment shapes

  • Local devcontextstore-server --config kv-service/configs/server-test.toml, backed by a local NVMe directory.
  • Docker / Compose — run make docker to build the image, then use kv-service/deploy/docker/docker-compose.yml; the Docker build includes the RDMA, io-uring, and metrics feature set.
  • Kuberneteskv-service/deploy/k8s/statefulset.yaml.
  • Systemdkv-service/deploy/systemd/contextstore-server.service.
  • JBOF over NVMe-oFkv-service/configs/server-nvmeof.toml plus the SPDK target and NVMe-oF initiator helpers in kv-service/deploy/jbof/. make build includes the pre-registered-slab, zero-memcpy RDMA WRITE data path.

See kv-service/configs/README.md for config fields and kv-service/deploy/README.md for the exact deployment commands.


Integration paths

Engine Connector class Module path
vLLM ContextStoreConnector contextstore.connector
Dynamo (KVBM) DynamoConnector contextstore.integrations.dynamo.connector
NIXL N/A (native plugin) Build nixl-plugin/contextstore/ — loads libcontextstore_nixl_client.so via CONTEXTSTORE backend

Roadmap

Directions we plan to work on next. This is a feature-level roadmap, not a release schedule; already-shipped capabilities are covered by the Architecture, KVService, Deployment, and Integration sections above.

Area Feature Status Notes
Connector Prefix-aware request routing planned Surface prefix-hit info back to upstream schedulers so they colocate requests that share a long prefix
Prefix Index Standalone Prefix Index Service planned Promote the in-process / Redis index into a dedicated service for cross-instance prefix discovery
Prefix Index Namespace and tenant isolation planned First-class model / tenant / namespace isolation, quotas and lifecycle rules
KVService Object lifecycle in progress TTL expiry, quota-based rejection, and cross-tier LRU eviction (L1 LRU is already live)
Data Path Raw JBOF block allocator planned Skip the filesystem: block allocator + placement policy directly against raw JBOF / NVMe-oF namespaces
Data Path GDS (GPUDirect Storage) in progress End-to-end libcufile + nvidia-fs path from NVMe straight into GPU HBM (requires datacenter-class GPUs)
Codec Higher-ratio KV codec planned INT4, per-layer quantization, and other layouts beyond the current INT8 codec

Contributions are welcome — please open an issue first to align on scope.


Development conventions

All code and doc changes must follow CLAUDE.md:

  • Python uses PEP 604 typing syntax (bytes | None, list[str]) with from __future__ import annotations at the top of every file; formatted with black, linted with ruff.
  • Rust is formatted with cargo fmt and must pass cargo clippy -- -D warnings; optional capabilities are gated behind Cargo features.
  • Storage access goes through the StorageBackend abstraction; hot paths avoid GIL-bound work and logging overhead.
  • The repository keeps exactly two top-level documents (README.md and CLAUDE.md). Do not add a docs/ tree — put design notes in code comments or PR descriptions.

License

Apache-2.0. See LICENSE.

About

Tiered shared KV cache storage for LLM inference. Extends GPU HBM to host memory, NVMe, and JBOF pools — enabling cross-instance prefix reuse for vLLM, Dynamo, and NIXL.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages