swap_blocks

Applicability

Note: For the Atlas A2 training product/Atlas A2 inference product, only the Atlas 800I A2 inference server and A200I A2 Box heterogeneous subrack are supported.

Function Description

Swaps in and out the cpu_cache and npu_cache.

For the swap-out function, the API starts four threads for concurrent execution. For the swap-in function, the API starts one D2D thread. To ensure stable performance, you are advised to bind processes to cores.

The swap-in function consists of two phases: H2D and D2D. To guarantee performance, the API allocates four block-sized buffers for pipeline copying. Therefore, you are advised to reserve sufficient device memory to prevent OOM errors.

Prototype

1

swap_blocks(src: KvCache, dst: KvCache, src_to_dst: Dict[int, int])

Parameters

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})

Returns

In normal cases, no value is returned.

If the input data type is incorrect and the source cache does not match the destination cache, a TypeError or ValueError is thrown.

If the input parameter is None, an AttributeError is thrown.

Constraints

None