KV Cache Eviction Policy

Problem

Evict least-recently-used KV blocks until at most max_blocks remain. Each block is [id, last_used_step]; keep the highest last_used values and, on ties, the larger block ids. Return the surviving ids sorted ascending.

The signature is kv_cache_eviction(blocks, max_blocks).

Hint

LRU eviction keeps the attention windows that recent tokens still need.

Test cases

case 1

in: [[[1,5],[2,9],[3,1]],2]

expected: [1,2]

case 2

in: [[[1,1],[2,2],[3,3]],2]

expected: [2,3]

case 3

in: [[[1,1]],1]

expected: [1]

case 4

in: [[[1,5],[2,5],[3,5]],2]

expected: [2,3]

Starter codePython

def kv_cache_eviction(blocks, max_blocks):
    # Your code here
    pass

4 test cases · run your code against them in the editor