note 03

Adding a KV-cache object

The first MiniCache reproduction was a tensor primitive. That was useful, but it still sat below the shape of an inference system.

The next step is a small cache object:

store per-layer keys and values, then apply the MiniCache pairwise path to a selected adjacent layer pair.

The implementation lives in:

minicache-pytorch

The new API

The new object is LayerKVCache. It stores key and value tensors with shape:

(layers, batch, sequence, hidden_dim)

The compression call is explicit about which adjacent layers are being merged:

import torch

from minicache_pytorch import LayerKVCache

keys = torch.randn(8, 2, 128, 64)
values = torch.randn(8, 2, 128, 64)
cache = LayerKVCache(keys=keys, values=values)

result = cache.compress_layer_pair(
    lower_layer=4,
    upper_layer=5,
    alpha=0.5,
    threshold=0.98,
)

compressed_cache = result.cache
retained_fraction = result.retained_token_fraction

The object does not mutate the original cache. It returns a new cache plus the key/value retention masks.

Why this is a better boundary

The primitive from the previous note answered:

can I merge two adjacent layer tensors?

The cache object asks a more useful systems question:

where would this sit in a decode-time cache path?

That boundary gives the repo a place to grow:

cache shape validation
layer-pair validation
memory accounting
retained-token reporting
future policy logic for selecting layer pairs

It is still small enough to read in one file.

Reproduce locally

Run the full test suite:

git clone https://github.com/rishabhsai/minicache-pytorch
cd minicache-pytorch
uv sync --extra dev
uv run --extra dev pytest

Current result:

15 passed

The cache-specific tests cover:

key/value shape validation
minimum cache rank validation
shape preservation after compression
only the upper layer in the selected pair is updated
retained tokens stay on the original upper-layer path
non-adjacent layer pairs fail clearly
cache byte accounting

Example output

The new example script creates a fake 8-layer cache, makes layers 4 and 5 partially similar, then compresses that pair:

uv run python examples/cache_object.py

Current local output:

cache shape: (8, 2, 128, 64)
cache bytes: 1,048,576
compressed pair: layers 4 -> 5
retained token fraction: 0.410

That retained fraction is the important signal. It means the cache path is not blindly compressing every token; the retention mask is visible enough to measure and debug.

What next

The next note should be a decode benchmark.

The benchmark should vary:

sequence length
hidden dimension
retention threshold
number of layer pairs compressed

The output should be a small table, not a wall of prose. The goal is to make the tradeoff visible: how much retention happens, how much extra compute the compression path adds, and where the primitive starts to look too expensive.