CopyKvCache

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

Copies the KV cache. D2D and D2H copy operations are supported.

To pipeline PullKvCache operations with other cache-consuming operations, you can allocate an additional intermediate cache. While other processes are using the current cache, you can prefetch the next cache block into this intermediate cache. After the current cache is released, copy the prefetched data from the intermediate cache to the designated location. This pipelines the PullKvCache latency, hiding it behind other operations and reducing the overall processing time.

In the common prefix scenario, before inference on a new request, the common prefix can be copied to a new memory space and merged with the KV cache of the current request for inference.

Prototype

1
2
3
4
5
6

Status CopyKvCache(const Cache &src_cache,
                   const Cache &dst_cache,
                   uint32_t src_batch_index = 0U,
                   uint32_t dst_batch_index = 0U,
                   uint64_t offset = 0U,
                   int64_t size = -1)

Parameters

Example

1

Status ret = llm_datadist.CopyKvCache(src_cache, dst_cache, 0, 0)

Returns

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

Constraints

Before calling this API, call the Initialize API to complete initialization. Only D2D and D2H copy operations are supported.