KV Cache Swapping

Function Description

KV cache swapping refers to swapping out the KV cache from the device to the host, or swapping it in from the host to the device.

In the PagedAttention scenario, due to the limited capacity of the KV cache, it is necessary to dynamically swap in and swap out the KV cache required by tasks based on factors such as task priority and execution status, thereby improving memory utilization.

Related APIs

Example

1
2
3
4
5
6
7
8

from llm_datadist import KVCache
...
npu_cache = kv_cache_manager.allocate_blocks_cache(npu_cache_desc, npu_cache_key)
cpu_cache = KvCache.create_cpu_cache(cpu_cache_desc, cpu_addrs) # cpu_addrs is obtained from the created CPU tensors.
# swap in
kv_cache_manager.swap_blocks(cpu_cache, npu_cache, {1:2, 3:4})
# swap out
kv_cache_manager.swap_blocks(npu_cache, cpu_cache, {1:2, 3:4})