wehub-resource-sync
143 lines
4.8 KiB
ReStructuredText
143 lines
4.8 KiB
ReStructuredText
KV Cache SDK
|
|
====================
|
|
|
|
The LMCache **SDK** lets you retrieve a request's KV cache from a LMCache
|
|
server, transform it on the CPU, and store it back. This can be used for KV
|
|
cache transformations, such as token dropping. In the example: we prefill a
|
|
batch of long prompts, drop half of each request's KV chunks, and show the
|
|
decode-throughput gain. The full runnable notebook lives at
|
|
`examples/token_dropping/token_dropping.ipynb
|
|
<https://github.com/LMCache/LMCache/tree/dev/examples/token_dropping>`_.
|
|
|
|
.. contents::
|
|
:local:
|
|
:depth: 2
|
|
|
|
|
|
Why KV Cache SDK
|
|
----------------
|
|
|
|
- **Improving Decode Throughput** when shrinking KV cache using token dropping.
|
|
Token dropping reduces the KV cache size, allowing more requests to fit in a
|
|
batch, improving decode throughput.
|
|
|
|
The SDK gives you the hooks to retrieve a request's KV, supply your own
|
|
function to edit the KV, and store the edited KV back. The SDK also provides a
|
|
batched-stream API to prefill, modify, and store the cache back before
|
|
decoding continues.
|
|
|
|
|
|
How it works
|
|
------------
|
|
|
|
A request flows through three phases on the batched-stream API:
|
|
|
|
- **prefill** — run each prompt through vLLM once (``max_tokens=1``); vLLM
|
|
computes the KV cache and stores it in LMCache.
|
|
- **modify** — the SDK retrieves the cached KV to CPU, hands it to your edit
|
|
function, and stores the result back.
|
|
- **decode** — continue generation against the smaller, edited cache.
|
|
|
|
The SDK runs on **CPU** and hands you KV tensors in ``HND`` order with shape
|
|
``[2, L, T, D]`` (K/V, layers, chunk-aligned tokens, ``num_kv_heads * head_dim``).
|
|
|
|
|
|
Configuration
|
|
-------------
|
|
|
|
To start the LMCache server with shared-memory transfer enabled, pass
|
|
``--shm-name`` and disable lazy L1 allocation with ``--no-l1-use-lazy``. If
|
|
shared memory is unavailable and these flags are not specified, the SDK falls
|
|
back to pickle.
|
|
|
|
.. code-block:: bash
|
|
|
|
lmcache server \
|
|
--l1-size-gb 150 \
|
|
--eviction-policy LRU \
|
|
--chunk-size 256 \
|
|
--port 6555 \
|
|
--http-port 8080 \
|
|
--shm-name lmcache_kvcache_sdk \
|
|
--no-l1-use-lazy
|
|
|
|
Then start vLLM with the LMCache MP connector.
|
|
|
|
.. code-block:: bash
|
|
|
|
vllm serve Qwen/Qwen3-8B \
|
|
--port 8000 \
|
|
--enforce-eager \
|
|
--gpu-memory-utilization 0.65 \
|
|
--kv-transfer-config '{
|
|
"kv_connector":"LMCacheMPConnector",
|
|
"kv_role":"kv_both",
|
|
"kv_connector_extra_config":{"lmcache.mp.port":6555}
|
|
}' \
|
|
--trust-remote-code \
|
|
--return-tokens-as-token-ids
|
|
|
|
The SDK keys the KV cache by token ids: ``create_request`` takes the prompt as
|
|
token ids, and every ``post_completion`` must report a ``token_id`` for each
|
|
generated token. The example gets these ids straight from vLLM by passing
|
|
``--return-tokens-as-token-ids``. Otherwise, if vLLM returns only text, the
|
|
``post_completion`` must tokenize each generated token back into a token id.
|
|
|
|
.. code-block:: python
|
|
|
|
import lmcache.sdk.kvcache as lmc_sdk
|
|
|
|
ctx = lmc_sdk.connect(
|
|
url="tcp://localhost:6555", # must match --port
|
|
http_url="http://localhost:8080", # must match --http-port
|
|
model_name="Qwen/Qwen3-8B",
|
|
timeout=60,
|
|
)
|
|
...
|
|
lmc_sdk.close(ctx)
|
|
|
|
|
|
Writing a custom edit function
|
|
------------------------------
|
|
|
|
An edit function takes the retrieved KV tensor and its token ids and returns the
|
|
edited ``(kv, tokens)``. ``batch.modify(fn)`` applies it to every stream.
|
|
|
|
``modify`` operates only on the **chunk-aligned** prefix. A trailing partial
|
|
chunk is tracked by the SDK and re-sent on the next ``decode``, so
|
|
``tokens`` arrives already truncated to the cached length.
|
|
|
|
|
|
API reference
|
|
-------------
|
|
|
|
.. list-table::
|
|
:header-rows: 1
|
|
:widths: 45 55
|
|
|
|
* - Function / method
|
|
- Description
|
|
* - ``lmc_sdk.connect(url, http_url, model_name, timeout=60.0)``
|
|
- Create an SDK context and register the transfer context, pass it to
|
|
every other call.
|
|
* - ``lmc_sdk.close(ctx)``
|
|
- Close the context and release resources. Called when done with the SDK.
|
|
* - ``lmc_stream.create_request(ctx, post_completion,
|
|
prompt_token_ids, cache_salt="")``
|
|
- Create one request stream to add to a batch.
|
|
* - ``lmc_batch.LMCacheBatchedStream()``
|
|
- Create an empty batch.
|
|
* - ``batch.add(stream)``
|
|
- Register a stream to the batch.
|
|
* - ``batch.prefill(sampling_params)``
|
|
- Prefill every stream once (``max_tokens`` forced to 1). Returns a
|
|
``Metrics`` report.
|
|
* - ``batch.modify(fn)``
|
|
- Apply the edit function ``fn`` to every stream's cached KV. Returns a
|
|
``Metrics`` report.
|
|
* - ``batch.decode(sampling_params)``
|
|
- Decode every stream. Returns a ``Metrics`` report.
|
|
|
|
```Metrics`` returns ``input_tokens``, ``input_tput`` for prefill, ``duration``
|
|
for modify, and ``output_tokens``, ``output_tput`` for decode.
|