CODING CHALLENGE · N°55

KV-Cache Eviction (Attention Sinks)

Medium AIInferenceCaching

An LLM’s KV cache grows with every token, so long chats need a way to drop old entries without wrecking quality. StreamingLLM’s insight: keep the first few tokens (the "attention sinks") plus a sliding window of the most recent ones, and throw away the middle. Compute which tokens survive. Solve it in Python or TypeScript, with hidden tests.

The problem

A sequence has n tokens (indices 0…n-1) in the KV cache. Under a StreamingLLM eviction policy you keep the first sink tokens (attention sinks) and the most recent window tokens, evicting everything in between. Implement kv_keep(n, sink, window): return the sorted list of token indices that remain in the cache.

EXAMPLE 1
Input n = 10, sink = 2, window = 3
Output [0, 1, 7, 8, 9]
first 2 sinks + last 3 recent; tokens 2–6 are evicted
EXAMPLE 2
Input n = 5, sink = 2, window = 3
Output [0, 1, 2, 3, 4]
the sink and window together already cover everything
CONSTRAINTS
  • Keep indices 0 … sink-1 (attention sinks) and n-window … n-1 (the recent window).
  • The two ranges may overlap or cover the whole sequence — return each surviving index once, sorted.
  • Clamp to the sequence bounds: if sink or window exceed n, just keep what exists.
SOLVE IT YOURSELF

Your turn — write it

Edit the stub, hit Run (or ⌘/Ctrl + Enter), and watch the hidden tests. Stuck? the hints are right above and Reveal solution is one click away.

YOUR TASK

Implement kv_keep(n, sink, window): form the set of sink indices and the set of recent-window indices, union them, and return the sorted result.

HINTS — 4 IDEAS
  1. Sink indices are range(min(sink, n)) — the first tokens, clamped to n.
  2. Recent-window indices are range(max(0, n - window), n) — the last window tokens.
  3. Union the two index sets so overlaps are not double-counted.
  4. Return them sorted. The evicted tokens are exactly the middle ones present in neither set.
CPython · WebAssembly

Explore the topic

See this challenge alongside everything else on the same subject — handbooks, system designs, algorithms and tools, in one place.