lmcache--lmcache/docs/design/sdk/kvcache.md

KV Cache SDK

Status: out-of-process, CPU-only client for the LMCache multiprocess (MP) server.

Goal

A small Python surface for moving KV cache tensors in and out of a running LMCache MP server, addressed by token ids — so an ML engineer can retrieve a prefix's KV, edit it, and store it back (e.g. token-dropping inside a KV connector).

import lmcache.sdk.kvcache as lmc_sdk

ctx = lmc_sdk.connect(
    url="tcp://localhost:5555",       # ZMQ url
    http_url="http://localhost:9000", # HTTP url for retrieving KV cache shape
    model_name="Qwen/Qwen3-8B",
)
kv = lmc_sdk.retrieve(ctx, tokens=[1, 2, 3, ...])  # [2, L, hit_tokens, D] or None
# ... edit kv ...
ok = lmc_sdk.store(ctx, kv=kv, tokens=[4, 5, 6, ...])
lmc_sdk.close(ctx)

See the token-dropping example.

The model layout must already be registered in the server by a vLLM instance that called REGISTER_KV_CACHE; the SDK reads that layout from /status and /config to configure itself.

Architecture

The SDK is a separate, CPU-only process from the LMCache server.

SDK process                                  LMCache MP server
-----------                                  -----------------
LMCacheKVCacheContext
  ├ ContiguousTransferWrapper
  │   └ EngineDrivenContext{Shm,Pickle} ──MQ──▶ EngineDrivenTransferModule
  │                                             └ StorageManager (L1 pool, locks, prefetch)
  └ MessageQueueClient ──────────────────ZMQ──▶ LookupModule (LOOKUP / QUERY_PREFETCH_STATUS)
  SharedMemory(name) ◀────────────────────────  L1 POSIX segment (SHM transport only)

• Control plane (ZMQ): lookup/prefetch, slot reservation, lock release, session end.
• Data plane: SHM when the server exposes an L1 pool, otherwise pickle over the MQ — both driven through one EngineDrivenContext, so the SDK never branches on transport.
• ContiguousTransferWrapper (wrapper/contiguous.py) bridges a contiguous [2, L, T, D] tensor to the per-chunk prepare/commit protocol and masks the SHM-vs-pickle difference.

Registration handshake

connect() builds the context, then register_kv_caches() runs once:

1. HTTP /config → chunk_size.
2. HTTP /status → the cache_context_meta entry for model_name: world_size and the GPU kv_cache_layout (num_layers, dtype, tokens_per_block, engine_kv_format, engine_kv_concrete_shape).
3. Decode geometry from the layout: num_kv_heads / block_size via the format-aware get_num_heads / get_block_size (on a meta probe tensor — no allocation), head_dim from the last dim; assert block_size == tokens_per_block.
4. Allocate a dummy 1-block HND buffer {layer.i: zeros([1, 2, NH, BS, HS], cpu)} — used only to register the layout; the data plane never touches it.
5. create_transfer_context(buffer) → EngineDrivenTransferContext.register(...), which sends REGISTER_KV_CACHE_ENGINE_DRIVEN_CONTEXT (server reserves the SHM pool / picks SHM vs pickle).
6. Wrap transfer_ctx.engine_driven_context in a ContiguousTransferWrapper.

instance_id is the SDK process PID.

Public API

Module functions in lmcache.sdk.kvcache; LMCacheKVCacheContext / KVCacheSDKError are exported from lmcache.sdk.

• connect(url, http_url, model_name, timeout=60.0) — open the MQ client, fetch config, run the handshake.
• retrieve(ctx, tokens, cache_salt="") → contiguous CPU [2, num_layers, hit_tokens, hidden_dim] for the cached prefix, or None (empty/sub-chunk input, or nothing cached).
• store(ctx, kv, tokens, cache_salt="") → bool. kv is [2, L, T, D]; len(tokens) must equal T; both are truncated to whole chunks before storing.
• close(ctx) — shut down the MQ client and ZMQ context.

Cache addressing

Both build an IPCCacheServerKey(model_name, world_size=1, worker_id=0, token_ids, start=0, end=<chunk-aligned>, request_id, cache_salt). The server resolves it to per-chunk ObjectKeys; cache identity = token-chunk hashes + model_name + kv_rank(worker_id) + cache_salt. request_id (store-/retrieve-<uuid>) only keys the per-request session, not cache identity.

• worker_id=0 is valid because the SDK is world_size == 1 (one shard per chunk).
• LOOKUP uses the worker_id=None (expand-to-all-workers) variant.

Flows

store — validate len(tokens) == kv.shape[2], truncate to whole chunks, then transfer_ctx.store(key, instance_id, kv_cpu): prepare_store returns writable SHM slots (filled in place) or None for pickle (gather all chunks); commit_store. Writes use mode "new", so already-cached chunks are deduplicated.

retrieve — Phase 0: LOOKUP + poll QUERY_PREFETCH_STATUS until non-None; if 0 → None. Phase 1: transfer_ctx.retrieve(key, instance_id): prepare_retrieve → torch.cat the chunk slots into one contiguous tensor → commit_retrieve. end_session always runs in finally.

Copy summary

Flow	SHM	Pickle
store	1 (tensor → SHM slot)	~2 (tensor → chunks, then serialize)
retrieve	1 (slot → contiguous)	~2 (deserialize, then → contiguous)

Constraints & known gaps

• world_size == 1 only, and a single non-hybrid kernel group only.
• Model must be pre-registered by a vLLM instance (REGISTER_KV_CACHE); the SDK reads the GPU layout from /status and cannot derive it from model_name alone.