KV Cache Management

Function Description

Before invoking LLM-DataDist to establish a link, you can allocate a memory block and register it with LLM-DataDist by calling the RegisterKvCache API. After the link is disconnected, call UnregisterKvCache to deregister it.

KV cache management involves the following APIs and functions.

API	Function
RegisterKvCache	Registers a local cache.
UnregisterKvCache	Deregisters a local cache.
PullKvCache	Pulls the cache from the remote node to the local cache. The call is not restricted by the node role.
PullKvBlocks	Pulls the cache from the remote node to the local cache by configuring a block list in the blocks cache transmission scenario. The call is not restricted by the node role.
PushKvCache	Pushes the cache from the local node to the remote node. The call is not restricted by the node role.
PushKvBlocks	Pushes the cache from the local node to the remote node by configuring a block list in the blocks cache transmission scenario. The call is not restricted by the node role.

Scenario

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 registration, deregistration, and transmission. Pseudocode examples are provided below for each role.

On the Prefill and Decode sides, Initialize the LLM-DataDist instance as described in Link Establishment, allocate memory, and register it.

void OnError(LlmDataDist &llm_data_dist, int64_t cache_id)
{
    if (cache_id > 0) {
        (void) llmDataDist.UnregisterKvCache(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};
std::vector<uint64_t> addrs;  // Populate as needed.
int64_t cache_id = -1;
auto ret = llm_data_dist.RegisterKvCache(kv_cache_desc, addrs, {}, cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] RegisterKvCache failed, ret = %u\n", ret);
    return -1;
}

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

The Prefill and Decode sides transmit the KV cache as needed. If the transmission fails, the corresponding resources must be released. There are two modes for KV cache transmission.

Call the PullKvCache API 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 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.
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;
}

When the service exits, perform link disconnection on the Prefill and Decode sides as described in Link Disconnection, call UnregisterKvCache to deregister the link, 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 registration 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.

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) and register it with the LLM-DataDist instance. 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, int64_t cache_id)
{
    if (cache_id > 0) {
        (void) llmDataDist.UnregisterKvCache(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};
std::vector<uint64_t> addrs;  // Populate as needed.
int64_t cache_id = -1;
auto ret = llm_data_dist.RegisterKvCache(kv_cache_desc, addrs, {}, cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] RegisterKvCache failed, ret = %u\n", ret);
    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 PullKvBlocks 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.
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;
}

When the service exits, perform link disconnection on the Prefill and Decode sides as described in Link Disconnection, call UnregisterKvCache to deregister the KV cache, and release the corresponding resources.

Exception Handling

LLM_KV_CACHE_NOT_EXIST occurs when the requested KV cache does not exist on the peer side. To address this, check whether the peer process is normal 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: Function Description