note 02

Reproducing MiniCache in PyTorch

MiniCache is useful to study because it compresses the KV cache along the depth axis instead of only asking which old tokens to keep.

The paper target is:

MiniCache: KV Cache Compression in Depth Dimension for Large Language Models

This reproduction does not claim paper parity. I am isolating the core tensor path first: decompose a pair of adjacent layer states, interpolate direction, preserve magnitude, and optionally retain tokens whose adjacent-layer similarity is too low to merge safely.

implementationminicache-pytorch

paperarXiv:2405.14366

scope tensor primitives, tests, benchmark

What is implemented

The first pass implements the smallest inspectable path:

magnitude / direction decomposition
adjacent-layer direction interpolation
reconstruction using the upper-layer magnitude
per-token cosine similarity between adjacent layer states
retention mask for low-similarity tokens
retention-aware adjacent-pair compression

That gives the repo a concrete API instead of only paper commentary.

import torch

from minicache_pytorch import compress_adjacent_pair_with_retention

lower = torch.randn(2, 16, 64)
upper = torch.randn(2, 16, 64)

compressed, keep_mask = compress_adjacent_pair_with_retention(
    lower,
    upper,
    alpha=0.5,
    threshold=0.92,
)

The merge path lives between neighboring layers: keep the magnitude stable, move the direction toward a shared representation.

Retention path

The retention step is the part that keeps this from being a blind average.

If adjacent layers are already similar for a token, the merge path is low risk. If they disagree sharply, that token stays on the original upper-layer path.

from minicache_pytorch import retention_mask_from_similarity

keep_mask = retention_mask_from_similarity(
    lower,
    upper,
    threshold=0.92,
)

The current implementation returns both the compressed tensor and the boolean mask, so the caller can inspect how much of the sequence was retained.

Tokens with unusually low cross-layer similarity are kept out of the merge path.

Reproduce locally

Clone the implementation repo and run the tests:

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

Current test coverage is focused on the low-level behavior:

normalization returns unit vectors
magnitude and direction reconstruct the original tensor
direction interpolation preserves unit norm
compression preserves shape
compression keeps the upper-layer magnitude
cosine similarity matches simple expected cases
low-similarity tokens are marked for retention
retained tokens stay on the original upper-layer path

First benchmark

The benchmark script measures the pairwise path on random tensors shaped like a small batch of cached states:

uv run python scripts/bench_pair.py

On my current local run, using tensors with shape (4, 512, 128), the output was:

plain compress_adjacent_pair: 0.279 ms
retention-aware pair path:   0.414 ms

This is not a model-quality benchmark. It is a sanity check for the primitive itself before wiring the method into a fuller cache-compression path.

What next

This first pass gives me a tested tensor primitive. The next notes should move from the primitive toward something closer to an inference system.

The sequence I want to build next:

Cache object integration
Wrap the pairwise primitive in a small KV-cache object so the code looks more like a decode-time path instead of a standalone tensor function.
Layer selection policy
Decide which adjacent layers are candidates for merging, and make that policy visible in the API.
Decode benchmark
Measure the compression path across longer sequence lengths and show the memory / latency tradeoff in a small table.
PrefixKV track
Start the next paper note around token retention for vision-language prompts, then define the smallest PyTorch primitive worth implementing.