Start from brute force
The direct version walks every index and keeps it if it is a sink or in the window: [i for i in range(n) if i < sink or i >= n - window]. It is correct, clamping included, because the test only ever runs on real indices, and it comes out sorted. But it costs O(n) time for O(k) output, where k is at most sink + window: a million-token stream with a 1,024-token budget scans a million indices to keep 1,024. The other quick version, list(range(sink)) + list(range(n - window, n)), is fast and wrong: it repeats overlapping indices and, for n = 3 with budgets of 10, invents tokens 3 to 9 and negative indices.
The approach
The survivors are two contiguous ranges, so build them from their bounds instead of testing every index. The reference kv_keep clamps each range to the sequence: sinks = set(range(min(sink, n))) holds the first tokens, and recent = set(range(max(0, n - window), n)) holds the last window. The min stops the sinks running past the end, and the max(0, …) stops the window starting at a negative index. sorted(sinks | recent) returns the union in order; the set union is what removes indices both ranges cover when sink + window reaches n.
The invariant: every index in the answer lies in [0, n) and in at least one range, and appears exactly once. The TypeScript kvKeep adds both ranges to one Set and sorts with (a, b) => a - b, because the default sort compares numbers as strings. Since both ranges are already sorted, concatenating range(0, a) and range(max(a, b), n), with a and b the clamped bounds, gives the same list in O(k) without a set.
Walk through an example
For the first example, n = 10, sink = 2, window = 3, the two ranges do not touch, so five tokens survive and the middle five are evicted.
| step | expression | value |
|---|
| 1 | sinks = set(range(min(2, 10))) | {0, 1} |
| 2 | recent = set(range(max(0, 10 - 3), 10)) | {7, 8, 9} |
| 3 | sinks | recent | {0, 1, 7, 8, 9} |
| 4 | sorted(...) | [0, 1, 7, 8, 9]; tokens 2 to 6 are in neither range |
Complexity
Time O(k log k) for the k ≤ sink + window survivors: building the two clamped sets is O(k) and sorting their union is O(k log k), with no dependence on n beyond two subtractions. Space O(k). The scan over every index costs O(n) time for the same O(k) output, and concatenating the two ranges brings the best case down to O(k).
Edge cases to test
kv_keep(5, 2, 3) → [0, 1, 2, 3, 4]. The ranges meet exactly, so nothing is evicted.kv_keep(6, 4, 4) → [0, 1, 2, 3, 4, 5]. Indices 2 and 3 are in both ranges; a list concatenation returns them twice.kv_keep(3, 10, 10) → [0, 1, 2]. Unclamped, range(3 - 10, 3) yields -7 to -1, and Python lists accept negative indices silently, so a lookup reads the wrong token instead of failing.kv_keep(4, 0, 2) → [2, 3]. No sinks is plain sliding-window attention, the setup StreamingLLM showed degrading once the first tokens leave the cache.kv_keep(10, 2, 0) → [0, 1]. A zero window keeps only the sinks, because range(10, 10) is empty.kv_keep(0, 4, 8) → []. An empty sequence keeps nothing.
Common mistakes
- Concatenating the two ranges without deduplicating — when
sink + window exceeds n they overlap and indices appear twice. - Not clamping:
range(n - window, n) starts at a negative index when window > n, and range(sink) invents tokens past the end when sink > n. - Off-by-one on the window start: the last
window tokens begin at n - window, not n - window + 1. - Dropping the sinks to save a few slots — a window-only cache is the configuration StreamingLLM showed degrading once the first tokens are evicted.
- In TypeScript, sorting numbers with the default comparator:
[10, 9].sort() compares strings and leaves [10, 9].
Follow-ups an interviewer asks next
- Tokens arrive one at a time. Maintain the cache incrementally.
- Keep the sinks fixed and hold the window in a ring buffer or deque of capacity
window. Each new token is appended; once the cache holds sink + window tokens, the oldest window entry, index t - window when token t arrives, is evicted. Each step is O(1), and no step recomputes the whole set. - What positions does the model see for the kept tokens?
- StreamingLLM assigns positions by place in the cache, not in the original text: with sinks 0 and 1 and window 7 to 9, the model sees positions 0 to 4. With rotary position embeddings that means caching keys before the rotation and applying it at each decoding step. Using the original indices would push positions past the lengths the model was trained on as the stream grows.
- Size the window from a memory budget in bytes.
- Each cached token stores one key and one value per layer, so it costs
2 × layers × kv_heads × head_dim × bytes_per_value. Divide the budget by that, subtract the sinks, and the rest is the window. Grouped-query attention reduces kv_heads and a lower-precision cache reduces bytes_per_value, which is why both stretch the window within the same memory. - Can you keep important middle tokens instead of dropping all of them?
- Yes, and that is the next family of policies. Heavy-hitter methods such as H2O keep the tokens that have accumulated the most attention so far, alongside the recent window. They need per-token attention statistics, so they cost more bookkeeping than a fixed sinks-plus-window rule, and they still cannot bring back a token once it has been evicted.
Where this shows up
The KV cache is the memory bottleneck of LLM serving: every generated token adds a key and a value per layer, so a long chat or a streaming agent eventually runs out of GPU memory. A sinks-plus-window policy, from the StreamingLLM paper, caps that memory at a fixed number of tokens while keeping generation stable, which suits always-on assistants and long agent sessions. It trades recall of the middle of a conversation for bounded memory, so it extends how long a model can run, not how much it remembers.
Reference solution
Reference solution — Python
def kv_keep(n, sink, window):
sinks = set(range(min(sink, n))) # first tokens (attention sinks)
recent = set(range(max(0, n - window), n)) # most recent tokens
return sorted(sinks | recent) # union, sorted
print(kv_keep(10, 2, 3)) # -> [0, 1, 7, 8, 9]
print(kv_keep(5, 2, 3)) # -> [0, 1, 2, 3, 4]
Reference solution — TypeScript
function kvKeep(n: number, sink: number, window: number): number[] {
const keep = new Set<number>();
for (let i = 0; i < Math.min(sink, n); i++) keep.add(i); // attention sinks
for (let i = Math.max(0, n - window); i < n; i++) keep.add(i); // recent window
return Array.from(keep).sort((a, b) => a - b);
}
console.log(kvKeep(10, 2, 3)); // -> [0, 1, 7, 8, 9]
console.log(kvKeep(5, 2, 3)); // -> [0, 1, 2, 3, 4]