swap_blocks

Applicability

Note: For the Atlas A2 training products/Atlas A2 inference products, only the Atlas 800I A2 inference server, Atlas 300I A2 inference card, and A200I A2 Box heterogeneous components are supported.

Function Description

Swaps in and out the CPU cache (cpu_cache) and NPU cache (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), a TypeError or ValueError is thrown.

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

Restrictions

This API can be used only in the PagedAttention scenario.