CopyKvBlocks

Applicability

Product	Supported (Yes/No)
Atlas 350 Accelerator Card	No
Atlas A3 training product/Atlas A3 inference product	Yes
Atlas A2 training product/Atlas A2 inference product	Yes
Atlas 200I/500 A2 inference product	No
Atlas inference product	No
Atlas training product	No

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

Copies the KV cache by specifying a block list in the PagedAttention scenario. D2D, D2H, and H2D copy operations are supported.

In the D2D scenario, when multiple responses need to share the same partially filled block, the new tokens must be copied to a new block for continued iteration.
H2D and D2H copy operations are mainly used for the swap-in and swap-out of cache memory at the corresponding block_index.

Prototype

Status CopyKvBlocks(const Cache &src_cache,
                    const Cache &dst_cache,
                    const std::vector<uint64_t> &src_blocks,
                    const std::vector<std::vector<uint64_t>> &dst_blocks_list)

Parameters

Parameter	Input/Output	Description
src_cache	Input	Source cache.
dst_cache	Input	Destination cache.
src_blocks	Input	Block index list of the source cache.
dst_blocks_list	Input	Block index list of the destination cache. A group of src_blocks can be copied to multiple groups of dst_blocks.

Example

Status ret = llm_datadist.CopyKvCache(src_cache, dst_cache, {1,2}, {{1,2},{3,4}})

Returns

LLM_SUCCESS: Success.
LLM_PARAM_INVALID: Incorrect parameter.
Other values: Failure.

Constraints

Before calling this API, call the Initialize API to complete initialization. H2H copy operations are not supported.

Parent topic: LlmDataDist APIs