KV Cache SDK#

LMCache SDK 允许您从 LMCache 服务器检索请求的 KV 缓存，在 CPU 上进行转换，然后再存储回去。这可以用于 KV 缓存转换，例如丢弃令牌。在这个例子中：我们预填充一批长提示，丢弃每个请求的一半 KV 块，并展示解码吞吐量的提升。完整的可运行笔记本位于 examples/token_dropping/token_dropping.ipynb。

为什么选择 KV Cache SDK#

  • 通过丢弃令牌 来提高解码吞吐量。当使用丢弃令牌缩小 KV Cache 时，KV Cache 的大小减少，使得更多请求可以适应一个批次，从而提高解码吞吐量。

SDK 提供了钩子来检索请求的 KV，允许您提供自己的函数来编辑 KV，并将编辑后的 KV 存储回去。SDK 还提供了一个批量流 API，以便在继续解码之前进行预填充、修改和存储缓存。

它是如何工作的#

请求在批处理流 API 上经历三个阶段：

  • 预取 — 将每个提示通过 vLLM 运行一次（max_tokens=1）；vLLM 计算 KV Cache 并将其存储在 LMCache 中。

  • 修改 — SDK 将缓存的 KV 取回到 CPU，交给你的编辑函数，并将结果存储回去。

  • 解码 — 在更小的、编辑过的缓存上继续生成。

SDK 在 CPU 上运行，并以 HND 顺序提供形状为 [2, L, T, D] 的 KV 张量（K/V、层、块对齐的标记、num_kv_heads * head_dim）。

配置#

要启动启用共享内存传输的 LMCache 服务器，请传递 --shm-name 并使用 --no-l1-use-lazy 禁用懒惰的 L1 分配。如果共享内存不可用且未指定这些标志，SDK 将回退到 pickle。

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

然后使用 LMCache MP 连接器启动 vLLM。

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

SDK通过令牌 ID 对 KV Cache 进行键控：create_request 将提示作为令牌 ID，且每个 post_completion 必须报告每个生成令牌的 token_id。该示例通过传递 --return-tokens-as-token-ids 直接从 vLLM 获取这些 ID。否则，如果 vLLM 仅返回文本，则 post_completion 必须将每个生成的令牌重新标记为令牌 ID。

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)

编写自定义编辑函数#

编辑函数接受检索到的 KV 张量及其令牌 ID，并返回编辑后的 (kv, tokens)。 batch.modify(fn) 将其应用于每个流。

modify 仅在 块对齐 前缀上操作。尾部的部分块由 SDK 跟踪，并在下一个 解码 时重新发送，因此 tokens 已经被截断到缓存的长度。

API 参考#

函数 / 方法

描述

lmc_sdk.connect(url, http_url, model_name, timeout=60.0)

创建一个 SDK 上下文并注册传输上下文，将其传递给每个其他调用。

lmc_sdk.close(ctx)

关闭上下文并释放资源。在完成 SDK 的使用时调用。

lmc_stream.create_request(ctx, post_completion, prompt_token_ids, cache_salt="")

创建一个请求流以添加到批处理中。

lmc_batch.LMCacheBatchedStream()

创建一个空批次。

batch.add(stream)

将流注册到批处理中。

batch.prefill(sampling_params)

一次 Prefill 每个流（max_tokens 强制为 1）。返回一个 Metrics 报告。

batch.modify(fn)

将编辑函数 fn 应用到每个流的缓存 KV。返回一个 Metrics 报告。

batch.decode(sampling_params)

解码每个流。返回一个 Metrics 报告。

`Metrics` 返回 input_tokens、input_tput 用于 Prefill，duration 用于修改，以及 output_tokens、output_tput 用于解码。