wehub-resource-sync
117 lines
5.6 KiB
Markdown
117 lines
5.6 KiB
Markdown
# 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).
|
|
|
|
```py
|
|
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](../../../examples/kvcache_sdk/e2e_kv_edit.py).
|
|
|
|
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](../../../lmcache/sdk/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 `ObjectKey`s;
|
|
**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.
|