KV compression - StreamingLLM, H2O, SnapKV

Open in Colab ▶️ Run this notebook in Colab

KV compression - StreamingLLM, H2O, SnapKV#

Track 05 - Serving · Notebook 03 · Runtime: ≈30 s on CPU

Prerequisites: 05_serving/02 (KV variants), 01_inference/01 (KV cache).

Papers:

  • Xiao et al. 2023, Efficient Streaming Language Models with Attention Sinks (StreamingLLM, 2309.17453).

  • Zhang et al. 2023, H2O: Heavy-Hitter Oracle (2306.14048).

  • Li et al. 2024, SnapKV: LLM Knows What You are Looking for Before Generation (2404.14469).

What#

Long contexts inflate the KV cache linearly. Three techniques reduce cached tokens (not bytes per token) during decoding:

  • StreamingLLM. Keep a small number of attention sinks at the very start of the context plus a sliding window at the end. Throw away the middle.

  • H2O. Score every token by its cumulative attention weight so far; evict the lowest-scoring tokens (“heavy hitters” stay).

  • SnapKV. At prompt-encoding time, look at the attention pattern of the final observation window; keep tokens that window attends to, drop the rest. Fixed cost, big savings.

We implement all three against a synthetic attention-score matrix and verify: (a) each policy keeps ≤ budget tokens; (b) H2O retains more ground-truth-important tokens than StreamingLLM; (c) SnapKV works in one pass at encoding time.

from llm_systems_cookbook.nb import bootstrap

import numpy as np

s = bootstrap("05_serving_03_kv_compression_streamingllm_h2o_snapkv")

Ground-truth attention matrix#

Simulate a T = 512-token sequence. The attention score matrix A[q, k] = the attention weight from query q to key k, rows normalised. We plant 32 “important” tokens scattered throughout that every query attends to heavily. A good compression keeps most of them.

set_seed(0)
rng = np.random.default_rng(0)
T = 512
IMPORTANT = sorted(rng.choice(T, size=32, replace=False).tolist())

# Base attention: uniform over earlier positions; boost "important" positions.
A = np.zeros((T, T), dtype=np.float32)
for q in range(T):
    for k in range(q + 1):
        A[q, k] = 1.0
    for k in IMPORTANT:
        if k <= q:
            A[q, k] += 8.0  # heavy-hitters
A += rng.exponential(0.1, size=A.shape) * np.tri(T, T)
# Row-normalise.
A = A / (A.sum(axis=1, keepdims=True) + 1e-12)
print(f"attention matrix {A.shape}, important tokens = {len(IMPORTANT)}")

Three compression policies#

Each returns a set of kept key indices under a budget B.

def streaming_llm(T: int, budget: int, n_sinks: int = 4) -> set[int]:
    '''n_sinks at the start + (budget - n_sinks) at the end.'''
    sinks = set(range(n_sinks))
    window = set(range(T - (budget - n_sinks), T))
    return sinks | window


def h2o(A: np.ndarray, budget: int) -> set[int]:
    '''Score each key by its cumulative attention across queries; keep top-budget.'''
    scores = A.sum(axis=0)
    top = np.argsort(-scores)[:budget]
    return set(top.tolist())


def snapkv(A: np.ndarray, budget: int, window: int = 32) -> set[int]:
    '''At encoding time, score keys by attention from the last ``window`` queries.'''
    T = A.shape[1]
    scores = A[-window:].sum(axis=0)
    top = np.argsort(-scores)[:budget]
    return set(top.tolist())


BUDGET = 64

kept_stream = streaming_llm(T, budget=BUDGET)
kept_h2o = h2o(A, budget=BUDGET)
kept_snap = snapkv(A, budget=BUDGET)

for name, kept in [("StreamingLLM", kept_stream), ("H2O", kept_h2o), ("SnapKV", kept_snap)]:
    recall = len(kept & set(IMPORTANT)) / len(IMPORTANT)
    print(f"  {name:<14}  kept={len(kept)}  important-token recall = {recall:.2%}")

recall_stream = len(kept_stream & set(IMPORTANT)) / len(IMPORTANT)
recall_h2o = len(kept_h2o & set(IMPORTANT)) / len(IMPORTANT)
recall_snap = len(kept_snap & set(IMPORTANT)) / len(IMPORTANT)

s.check("streaming_respects_budget", lambda: len(kept_stream) <= BUDGET, msg=f"{len(kept_stream)}")
s.check("h2o_respects_budget",       lambda: len(kept_h2o) <= BUDGET,    msg=f"{len(kept_h2o)}")
s.check("snapkv_respects_budget",    lambda: len(kept_snap) <= BUDGET,   msg=f"{len(kept_snap)}")

s.check(
    "h2o_best_recall_of_important_tokens",
    lambda: recall_h2o >= recall_stream,
    msg=f"h2o={recall_h2o:.1%}  streaming={recall_stream:.1%}",
)
s.check(
    "snapkv_comparable_to_h2o",
    lambda: recall_snap >= recall_h2o - 0.15,
    msg=f"snapkv={recall_snap:.1%}  h2o={recall_h2o:.1%}",
)
s.check(
    "all_policies_nontrivially_better_than_random",
    lambda: min(recall_h2o, recall_snap) > BUDGET / T + 0.1,
    msg=f"random baseline ~{BUDGET / T:.1%}",
)

Compression vs quality on one plot#

x-axis: how much of the cache was thrown away (higher = more aggressive). y-axis: recall of the planted “important” tokens (higher = better quality proxy). The three methods sit at the same compression level here but separate cleanly on quality - H2O and SnapKV find the heavy hitters, StreamingLLM keeps what’s nearby.

import matplotlib.pyplot as plt

methods = [
    ("StreamingLLM", kept_stream, recall_stream, "tab:blue"),
    ("H2O",          kept_h2o,    recall_h2o,    "tab:orange"),
    ("SnapKV",       kept_snap,   recall_snap,   "tab:green"),
]
random_recall = BUDGET / T

fig, ax = plt.subplots(figsize=(6.8, 4.2))
for name, kept, recall, color in methods:
    compression = 1.0 - len(kept) / T
    ax.scatter(compression, recall, s=220, color=color, edgecolor="black",
               lw=0.8, zorder=5, label=name)
    ax.annotate(name, (compression, recall), xytext=(8, 6),
                textcoords="offset points", fontsize=10)

ax.axhline(random_recall, color="tab:gray", linestyle="--",
           label=f"random baseline ({random_recall:.1%})")
ax.axhline(1.0, color="black", linestyle=":", alpha=0.4, label="oracle")
ax.set_xlabel("compression ratio (fraction of KV evicted)")
ax.set_ylabel("important-token recall")
ax.set_title(f"KV compression Pareto @ budget={BUDGET}/{T} tokens")
ax.set_xlim(0, 1)
ax.set_ylim(0, 1.05)
ax.grid(True, alpha=0.3)
ax.legend(loc="lower left", fontsize=9)
plt.tight_layout()
plt.show()

Exercises#

  1. Sweep budget. Plot important-token recall vs budget for all three policies. StreamingLLM has a step function (fixed sinks + fixed window); H2O/SnapKV have smooth curves.

  2. Combine. Union a StreamingLLM window with H2O heavy hitters. Typical real systems do this - the window covers recent interactions, H2O catches distant must-keep tokens.

  3. Real integration. Plug a compression policy into 01_inference/01’s KV loop - on each decode step, evict tokens not in kept.

References#

  • StreamingLLM, H2O, SnapKV papers for the three algorithms.

  • KIVI (notebook 04) attacks the problem from a different angle: keep all tokens, but quantise K/V to 2 bits.

s.summary()
s.save()