KV Cache Management

Function Description

During LLM-DataDist initialization, a memory pool of a specified size is allocated in advance. The size is specified by OPTION_BUF_POOL_CFG. The subsequent KV cache allocation and deallocation are performed on the memory pool, which saves time compared with allocating a memory block each time.

KV cache management involves the following APIs and functions.

API	Description
AllocateCache	Allocates a cache.
DeallocateCache	Deallocates a cache.
PullKvCache	Pulls the cache from the remote node to the local cache. This API can be called only when the role is Decoder.
PullKvBlocks	Pulls the cache from the remote node to the local cache by configuring a block list in the blocks cache transmission scenario. This API can be called only when the role is Decoder.
CopyKvCache	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.
CopyKvBlocks	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.
PushKvCache	Pushes the cache from the local node to the remote node. This API can be called only when the role is Prompt.
PushKvBlocks	Pushes the cache from the local node to the remote node by configuring a block list in the blocks cache transmission scenario. This API can be called only when the role is Prompt.

Application Scenarios

It is mainly used for KV cache transmission and movement between distributed clusters.

Example (Typical Cache Transmission Scenario)

This example demonstrates API usage in a typical cache transmission scenario, covering KV cache allocation, deallocation, and transmission. Pseudocode examples are provided below for each role.

On the Prefill and Decode sides, initialize the LLM-DataDist instance and establish a link as described in Link Establishment.

For each request, allocate a KV cache of the required size on both sides. If allocation fails, deallocate any previously allocated resources.

void OnError(LlmDataDist &llm_data_dist, Cache &cache)
{
    if (cache.cache_id > 0) {
        (void) llmDataDist.DeallocateCache(cache.cache_id);
    }
    llm_data_dist.Finalize();
}
CacheDesc kv_cache_desc{};
kv_cache_desc.num_tensors = NUM_TENSORS;
kv_cache_desc.data_type = DT_INT32;
kv_cache_desc.shape = {8, 16};
Cache cache{};
auto ret = llm_data_dist.AllocateCache(kv_cache_desc, cache);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] AllocateCache failed, ret = %u\n", ret);
    OnError(llm_data_dist, cache);
    return -1;
}

Transfer the KV cache from the Prefill side to the Decode side in either of the following ways:

Call the PullKvCache API on the Decode side to pull the KV cache.

// Perform full inference on the Prefill side, write the results to the cache, and notify the Decode side that the cache is ready for pulling.
// Pull the cache to the Decode side.
CacheIndex cache_index{PROMPT_CLUSTER_ID, 1, 0};
ret = llm_data_dist.PullKvCache(cache_index, cache, 0);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PullKvCache failed, ret = %u\n", ret);
    return -1;
}
// Perform incremental inference.

Call the PushKvCache API on the Prefill side to push the KV cache.

// After the Prefill side completes the first-layer computation, it can immediately push the data in the transmission thread to overlap most of the transmission time with the computation time.
CacheIndex dst_cache_index{DECODE_CLUSTER_ID, 1, 0};
KvCacheExtParam ext_param{};
ext_param.src_layer_range =  std::pair<int32_t, int32_t>(0, 0);
ext_param.dst_layer_range =  std::pair<int32_t, int32_t>(0, 0);
// Number of tensors at each layer. Change it based on the actual model.
ext_param.tensor_num_per_layer = 2;
ret = llm_data_dist.PushKvCache(cache, dst_cache_index, 0, -1, ext_param);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PushKvCache failed, ret = %u\n", ret);
    return -1;
}

Release the KV cache for the corresponding requests on the Prefill and Decode sides, based on the cache usage timing in the service.

ret = llm_data_dist.DeallocateCache(cache.cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] DeallocateCache failed, ret = %u\n", ret);
} else {
    printf("[INFO] DeallocateCache success\n");
} 

When the service exits, disconnect the links on Prefill and Decode sides as described in Link Disconnection and call finalize to release resources.

Example (Typical Blocks Cache Transmission Scenario)

This example demonstrates API usage in a blocks cache transmission scenario, covering KV cache allocation, deallocation, and transmission. Pseudocode examples are provided below for each role.

Initialize the LLM-DataDist instance and establish a link on the Prefill and Decode sides by following the example in Link Establishment.

Call AllocateCache to allocate the KV cache at each layer of the model on the Prefill and Decode sides based on the calculated number of blocks (num_block). The upper-layer framework reuses the created KV cache (of size num_block) for different requests, and deallocates the allocated memory after the service process ends.

void OnError(LlmDataDist &llm_data_dist, Cache &cache)
{
    if (cache.cache_id > 0) {
        (void) llm_data_dist.DeallocateCache(cache.cache_id);
    }
    llm_data_dist.Finalize();
}
CacheDesc kv_cache_desc{};
kv_cache_desc.num_tensors = NUM_TENSORS;
kv_cache_desc.data_type = DT_INT32;
kv_cache_desc.shape = {8, 16};
Cache cache{};
auto ret = llm_data_dist.AllocateCache(kv_cache_desc, cache);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] AllocateCache failed, ret = %u\n", ret);
    OnError(llm_data_dist, cache);
    return -1;
}

When a new request arrives on the Prefill side, the inference framework allocates a block_index for that request. Once model inference finishes, the KV cache associated with the request resides in the memory region identified by the allocated block_index. The model output and the corresponding block_table are then passed as input to the Decode-side inference model.

When a new request arrives on the Decode side, the inference framework also allocates a block_index for that request. It then calls the pull_blocks API to transmit the KV cache to the designated location, based on the mapping between the Prefill-side block_index and the Decode-side block_index.

Call the PullKvBlocks API on the Decode side to pull the KV cache.

// Perform full inference on the Prefill side, write the results to the cache, and notify the Decode side that the cache is ready for pulling.
// Pull the cache to the Decode side.
CacheIndex cache_index{PROMPT_CLUSTER_ID, 1, 0};
std::vector<uint64_t> prompt_blocks = {0, 1, 2, 3};
std::vector<uint64_t> decoder_blocks = {3, 2, 1, 0};
auto ret = llm_data_dist.PullKvBlocks(cache_index, cache, prompt_blocks, decoder_blocks);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PullKvBlocks failed, ret = %u\n", ret);
    return -1;
}
// Perform incremental inference.

Call the PushKvBlocks API on the Prefill side to push the KV cache.

// After the Prefill side completes the first-layer computation, it can immediately push the data in the transmission thread to overlap most of the transmission time with the computation time.
CacheIndex dst_cache_index{DECODE_CLUSTER_ID, 1};
KvCacheExtParam ext_param{};
ext_param.src_layer_range =  std::pair<int32_t, int32_t>(0, 0);
ext_param.dst_layer_range =  std::pair<int32_t, int32_t>(0, 0);
// Number of tensors at each layer. Change it based on the actual model.
ext_param.tensor_num_per_layer = 2;
std::vector<uint64_t> prompt_blocks = {0, 1, 2, 3};
std::vector<uint64_t> decoder_blocks = {3, 2, 1, 0};
ret = llm_data_dist.PushKvBlocks(cache, dst_cache_index, prompt_blocks, decoder_blocks, ext_param);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PushKvBlocks failed, ret = %u\n", ret);
    return -1;
}

Release the KV cache for the corresponding requests on the Prefill and Decode sides, based on the cache usage timing in the service.

ret = llm_data_dist.DeallocateCache(cache.cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] DeallocateCache failed, ret = %u\n", ret);
} else {
    printf("[INFO] DeallocateCache success\n");
}

When the service exits, disconnect the links on Prefill and Decode sides as described in Link Disconnection and release resources.

Troubleshooting

LLM_DEVICE_OUT_OF_MEMORY indicates insufficient device memory for KV cache allocation. To address this, check the OPTION_BUF_POOL_CFG size configured during initialization and the requested KV cache size. Also verify that all previously pulled KV caches have been properly released.
LLM_KV_CACHE_NOT_EXIST occurs when the requested KV cache does not exist on the peer side. To address this, check if the peer process is abnormal and whether inference for the request corresponding to the KV cache has been completed. This error does not affect the processing flow of other requests. You may retry the failed request after verification.
LLM_WAIT_PROCESS_TIMEOUT or LLM_TIMEOUT occurs when a KV cache pull operation times out, indicating a link issue. To address this, try disconnecting and re-establishing the link.
LLM_NOT_YET_LINK occurs when no link has been established with the remote cluster.

Parent topic: Functions