swap_blocks

Applicability

Note: For 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_cache: Cache, dst_cache: Cache, src_to_dst: Dict[int, int])

Parameters

Example

1
2
3
4
5
6
7

from llm_datadist import Cache
npu_cache = cache_manager.allocate_blocks_cache(npu_cache_desc, npu_cache_key)
cpu_cache = Cache.create_cpu_cache(cpu_cache_desc, cpu_addrs) # cpu_addrs is obtained from the created CPU tensors.
# swap in
cache_manager.swap_blocks(cpu_cache, npu_cache, {1:2, 3:4})
# swap out
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 (src and dst do not match), the TypeError or ValueError exception is reported.

If the input parameter is None, the AttributeError exception is reported.

Restrictions

This API can be used only in the PagedAttention scenario.