From Research Paper to Real Workload: KV-Cache Editing with the LMCache SDK

Calling ML engineers and researchers: benchmark your KV-cache editing ideas with LMCache! 🔬

Large language models are powerful, and as Jensen Huang pointed out at GTC 2026, KV cache is the key to making large language models run fast.

There are many different ways to optimize KV cache, such as token dropping and KV cache training. These deliver amazing decoding throughput improvement (1.7x when dropping 50% tokens) and accuracy improvement (4% when dropping 50% tokens). However, none of them land in the production-level stack.

Here comes LMCache SDK – a new SDK that helps you integrate your KV cache research with only a Jupyter notebook!

The challenge of landing KV cache optimization in production stack

Some KV-cache techniques can operate using only Key and Value tensors.

More advanced techniques, however, often depend on temporary intermediate tensors produced during the model’s forward pass. One important example is the query tensor, which is used to rank cached tokens to retain the most useful and remove tokens with less contribution.

There are two major practical barriers between a research prototype and a production-style inference benchmark.

1. Intermediate tensors are temporary

Query tensors are produced inside the attention computation and normally disappear after the operation completes. They are not model parameters, nor are they typically exposed through an inference engine’s public APIs. The LMCache SDK captures this using a small vLLM patch that forwards Query tensors as intermediate tensors through the KV connector.

2. Experimental work can interfere with inference

Running KV cache-ranking, compaction, or selection logic on the same GPU as inference can compete for compute cycles, memory bandwidth, and temporary memory.

When this logic is inserted directly into the inference hot path, resources that could otherwise serve requests are diverted to cache optimization. The potential benefit of reducing the KV cache then becomes tightly coupled with the cost of performing the optimization itself.

This makes results harder to interpret:

• Did the smaller cache improve batching?
• Did the editing overhead cancel the throughput gain?
• Is the algorithm effective but poorly colocated?
• Would it perform better on a CPU or another accelerator?

A useful research interface should therefore separate cache-editing logic from the inference engine as much as possible while still supporting end-to-end evaluation.

Introducing the LMCache SDK

The LMCache SDK provides a practical playground for KV-cache research.

It gives user applications hooks to:

• retrieve a request’s KV cache;
• retrieve supported intermediate tensors, currently including Query tensors;
• pass those tensors to a user-defined editing function;
• store the modified cache back into LMCache; and
• continue decoding using the edited cache.

Instead of implementing an algorithm around inference-engine scheduling, batching, cache layouts, and request management, researchers can write a function as though they were transforming a single request.

Researchers provide the transformation; LMCache handles the cache and the request plumbing.

At a high level, the editing function receives a mapping of available cache tensors together with the request’s token sequence:

def edit_kv(
    caches: Mapping[CacheKind, torch.Tensor],
    tokens: list[int],
):
    # Rank, drop, merge, compress, quantize,
    # or otherwise modify the cache.
    return edited_kv, edited_tokens

The SDK then applies this transformation across the selected request streams.

A three-phase experimental workflow

The batched-stream API organizes an experiment into three main phases.

Phase 1: Prefill

Each prompt is first sent through the inference engine. The SDK sends the requests to the inference engine to generate 1 token each. As a result, the engine constructs the requests’ KV cache and Query tensors, then stores it in LMCache.

Phase 2: Modify

Before normal decoding continues, the SDK retrieves the stored cache and passes it to the user’s custom editing function which takes the tensors and the tokens as the input, then returns the edited KV and tokens.

The function may rank and remove token positions, compress or quantize cached tensors, merge similar cache entries, apply layer- or head-specific cache budgets, use Query tensors to guide cache selection, etc.

The returned KV cache is then stored back to LMCache by the SDK.

Cache retrieval and editing are performed on the CPU, outside the primary inference-GPU execution path. This helps decouple the editing workload from inference scheduling and compute resources.

Phase 3: Decode

The requests resume generation using the modified cache.

Write for one request, run across many requests

Real serving systems process many concurrent requests, not just one prompt.

The SDK’s stream abstraction represents the lifecycle of an individual request. Multiple streams can then be grouped into an LMCacheBatchedStream and executed together. Researchers do not need to build a separate batched implementation to evaluate an algorithm under concurrent load.

The SDK coordinates:

• the participating request streams;
• concurrent cache retrieval;
• invocation of the editing function;
• storage of edited caches;
• continuation of generation; and
• collection of serving metrics.

This is especially handy for techniques where main systems benefit only appears when the smaller per-request footprint allows the inference engine to decode more requests concurrently under high load, such as token dropping.

Behind the scenes: from engine tensors to SDK tensors

LMCache internally stores and transfers caches in chunks, but the SDK hides most of that representation from the user.

The CPU-side KV tensor is presented in HND order with the [2, L, T, D] shape, where:

• 2 represents Key and Value;
• L is the number of layers;
• T is the chunk-aligned token dimension; and
• D is num_kv_heads × head_dim.

This allows editing functions to operate on a contiguous, request-oriented tensor. LMCache manages chunked storage and transport internally, while the researcher can work with contiguous request-level tensors.

Capturing Query tensors with the Q ring buffer

As previously explained, Query tensors are ephemeral, produced layer by layer during each forward pass. LMCache captures them using a temporary Q ring buffer inside the vLLM process.

For each forward step, the system allocates ring-buffer slots for the relevant layers and scheduled requests. As attention executes, Query tensors are staged in the ring buffer. When KV-cache offloading to LMCache is triggered, the corresponding Query tensors are transferred alongside the KV cache.

This gives researchers access to Query tensors after the forward pass has completed, without requiring the custom editing algorithm to run directly inside the attention computation.

After being copied to the Q ring buffer that acts as a staging buffer, the SDK then uses the existing KV cache store machinery to store the query tensor chunks under a suffixed model name. This allows LMCache to distinguish Query tensors from ordinary KV entries while reusing the surrounding cache-server, transport, and lookup infrastructure.

What can you build with the SDK?

Although the example only demonstrates token dropping, you can use the SDK for various ideas, including but not limited to: query-aware token selection, layer/head-specific cache budgets, and intermediate tensor cache management. If you have an idea for how KV caches can be edited, we’d love to see what you build!

LMCache SDK as a systems playground

The LMCache SDK enables faster research prototyping on KV cache editing algorithms. One algorithm may reduce memory but hurt generation quality. Another may preserve quality but take too long to execute. A third may improve throughput only under particular batching or memory-pressure conditions. That is precisely why a systems-level playground is useful.

Researchers can evaluate the complete trade-off across:

• cache-size reduction;
• editing overhead;
• tensor-transfer overhead;
• decode throughput, and
• generation quality.

Rather than stopping at theoretical memory savings or isolated attention simulations, an algorithm can be connected to a real inference workload and benchmarked under concurrent serving conditions.

How to get started

The repository provides runnable token-dropping examples:

• Random token dropping: requires only the KV cache.
• SnapKV token dropping: uses Query tensors for token-importance scoring.
• Google Colab SnapKV example: a smaller demonstration designed for a T4 GPU.

Token-dropping examples: https://github.com/LMCache/LMCache/tree/dev/examples/token_dropping

LMCache SDK documentation: https://docs.lmcache.ai/mp/sdk.html

Google Colab examples: https://drive.google.com/drive/folders/1ILctdh_Lf51qDUf1v00osRfoylfBIlOD?usp=sharing

Writing and benchmarking your own KV-cache editing algorithm now requires editing only one Jupyter Notebook cell!