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 ge.flowGraphMemMaxSize. 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.

**Table 1** Main KV cache management APIs and their functions
API	Description
allocate_cache	Allocates a cache. After successful allocation, the cache is simultaneously referenced by cache_id and cache_keys. The resources occupied by the cache are not actually released until all these references are removed. cache_keys is passed in when LLMRole is set to PROMPT.
deallocate_cache	Deallocates a cache. If the cache was associated with cache keys during memory allocation, actual deallocation is deferred until all cache keys are released. The cache key references can be removed in the following ways: The Decode side successfully calls the pull_cache API. The Prompt side calls the remove_cache_key API.
remove_cache_key	Removes a cache key. After a cache key is removed, the corresponding cache can no longer be pulled by pull_cache. This API can only be called when LLMRole is PROMPT. It does not need to be called when LLMRole is DECODER, as cache_key is not passed in during cache allocation (allocate_cache).
pull_cache	Pulls the KV cache from the corresponding Prompt node to the local KV cache based on CacheKey. This API can be called only when LLMRole is set to DECODER. CacheKey must be the same as that in allocate_cache.
copy_cache	Copies the KV cache. To pipeline pull_cache with other operations using the KV cache, an additional intermediate cache can be allocated. While other processes are using the KV cache, the KV cache required for the next step can be pulled into the intermediate cache. After the other processes finish using the current KV cache, the pulled cache can be copied to the designated location. This hides the latency of pull_cache, 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.
allocate_blocks_cache	Allocates the cache for multiple blocks in the PagedAttention scenario.
pull_blocks	Pulls a KV cache in the PagedAttention scenario. Unlike pull_cache, pull_blocks pulls the KV cache at the corresponding positions based on block_index.
copy_blocks	Copies a KV cache in the PagedAttention scenario. Unlike copy_cache, copy_blocks copies the KV cache at the corresponding positions based on block_index. There are also differences in application scenarios. copy_blocks is used when multiple responses need to share the same partially filled block. In this case, the new tokens must be copied to a new block for continued iteration.
swap_blocks	Swaps in or swaps out the KV memory corresponding to a block_index in the PagedAttention scenario. This API is used when users need to manage the KV memory by themselves.
transfer_cache_async	Transmits the cache data to the Decode side. This process is initiated by the Prompt node.
push_blocks	Pushes KVs from the corresponding Prompt node to the remote KV cache in the PagedAttention scenario.
push_cache	Pushes KVs from the corresponding Prompt node to the remote KV cache.

Application Scenarios

This capability is primarily used for KV cache transmission 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.

Initialize the LLM-DataDist instance and establish a link on the Prefill and Decode sides as described in Link Management.

On the Prefill side, allocate KV cache memory of the corresponding size for each request. Taking PyTorch as an example, convert the KV cache into torch tensors and perform full-model inference.

import torchair
import torch
import torch_npu
# Obtain kv_cache_manager from the initialized llm_datadist.
kv_cache_manager = llm_datadist.kv_cache_manager
# Create a CacheDesc based on the shape and total number of KV caches in the model. The shape provided here is for reference only; replace it with the actual KV cache shape from the network.
cache_desc = CacheDesc(num_tensors=4, shape=[4, 4, 8], data_type=DataType.DT_FLOAT16)
# Create cache_keys for the corresponding requests based on cluster_id specified during llm_datadist initialization. When the Prefill side uses a multi-batch model, create cache_keys equal to the batch size.
batch_num = 8
kv_cache = kv_cache_manager.allocate_cache(cache_desc, [CacheKey(prompt_cluster_id=0, req_id=i, model_id=0) for i in range(batch_num)])

# Convert the allocated KV cache into torch tensors.
kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0]
kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs)
# Pass the converted kv_tensors to the model for inference computation to generate the KV cache, and transmit the model output to the incremental inference model as the input.

Allocate a KV cache for model execution on the Decode side.

# Obtain kv_cache_manager from the initialized llm_datadist.
kv_cache_manager = llm_datadist.kv_cache_manager
# Create a CacheDesc based on the shape and total number of KV caches in the model.
cache_desc = CacheDesc(num_tensors=4, shape=[4, 4, 8], data_type=DataType.DT_FLOAT16)
# Call the allocate_cache API to allocate the KV cache memory for the corresponding request.
kv_cache = kv_cache_manager.allocate_cache(cache_desc)

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

On the Decode side, call the pull_cache API to pull the KV cache for the corresponding request into the allocated memory.

# Create the same cache_key used when allocating the cache on the Prefill side to pull the corresponding KV cache.
cache_key = CacheKey(prompt_cluster_id=0, req_id=1, model_id=0)
kv_cache_manager.pull_cache(cache_key, kv_cache, batch_index=1) # Pull the cache to the position with batch index 1.

On the Prefill side, call the transfer_cache_async API to transmit data to the Decode side.

from llm_datadist import LayerSynchronizer, TransferConfig


class LayerSynchronizerImpl(LayerSynchronizer):
    def __init__(self, events):
        self._events = events

    def synchronize_layer(self, layer_index: int, timeout_in_millis: Optional[int]) -> bool:
        self._events[layer_index].wait()
        return True


events = [torch.npu.Event() for _ in range(cache_desc.num_tensors // 2)]
# Execute the model. The model calls events[layer_index].record() after the computation of each layer is complete.
# Model execution is implemented by the user.
# user_model.Predict(kv_tensors, events)

# After the model execution is initiated, call transfer_cache_async to transmit data. The memory addresses of the KV cache tensors for each layer allocated by the Decode side need to be provided here.
transfer_config = TransferConfig(DECODER_CLUSTER_ID, decoder_kv_cache_addrs)
cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(events), [transfer_config])
# Wait for the transmission result synchronously.
cache_task.synchronize()

Taking PyTorch as an example, convert the KV cache into torch tensors and perform incremental model inference.

# Convert the allocated KV cache to the KV cache type required by the framework. Different frameworks need to provide an API to create a KV cache of the corresponding type based on the KV cache address. Taking PyTorch as an example:
# Perform the conversion and pull operations in any sequence.
kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0]
kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs)
# Split the converted tensors into the KV pairing format required by the framework. KV can be combined arbitrarily.
mid = len(kv_tensors) // 2
k_tensors = kv_tensors[: mid]
v_tensors = kv_tensors[mid:]
kv_cache_tensors = list(zip(k_tensors, v_tensors))

# Pass the converted kv_tensors to the model for iterative inference.
# Wait until the request for incremental inference is complete.

Release the KV cache memory for the corresponding request at the appropriate time based on the service logic.

# On the Prefill side, the actual memory release by deallocate_cache occurs only after the KV caches for all requests in a batch have been pulled. If the KV cache has not been pulled by the Decode side, remove_cache_key also needs to be called.
kv_cache_manager.remove_cache_key(cache_key_0)
kv_cache_manager.remove_cache_key(cache_key_2)
kv_cache_manager.deallocate_cache(kv_cache)

# On the Decode side, as cache_key is not required during allocation, only the deallocate API needs to be called for release.
kv_cache_manager.deallocate_cache(kv_cache)

When the service exits, disconnect the links on Prefill and Decode sides as described in Link Management 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 as described in Link Management.

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_blocks). In the Blocks Cache scenario, different requests reuse the allocated KV cache of the num_blocks size. The upper-layer framework manages this process, and the allocated memory is released when the service ends.

# Prefill
# Obtain kv_cache_manager from the initialized llm_datadist.
kv_cache_manager = llm_datadist.kv_cache_manager
# Create a CacheDesc based on the shape and total number of KV caches in the model. In the PagedAttention scenario, the KV cache shape is usually [num_blocks, block_size,...,...].
num_blocks = 10
block_mem_size = 128
cache_desc = CacheDesc(num_tensors=4, shape=[num_blocks, block_mem_size], data_type=DataType.DT_FLOAT16)
# Create a BlocksCacheKey for the corresponding request based on cluster_id used during llm_datadist initialization.
cache_key = BlocksCacheKey(prompt_cluster_id=0, model_id=0)
# Call the allocate_blocks_cache API to allocate KV cache memory.
kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc, cache_key)

# Convert the allocated KV cache to the KV cache type required by the framework. Different frameworks need to provide an API to create a KV cache of the corresponding type based on the KV cache address. Taking PyTorch as an example:
kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0]
kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs)

# Decode
kv_cache_manager = llm_datadist.kv_cache_manager
num_blocks = 10
block_mem_size = 128
cache_desc = CacheDesc(num_tensors=4, shape=[10, 128], data_type=DataType.DT_FLOAT16)
kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc)

kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0]
kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs)

When a new request arrives on the Prefill side, the inference framework allocates a corresponding 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. In this case, there are two methods:

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

# The Decode side adds a new request based on the information passed from the Prefill side and allocates the corresponding block_table.
# The Decode side pulls the KV cache to the corresponding block based on src_block_table of the incoming request and the newly allocated dst_block_table.
cache_key = BlocksCacheKey(prompt_cluster_id=0, model_id=0) # Input parameter when allocate_blocks_cache is called on the Prefill side.
kv_cache_manager.pull_blocks(cache_key, cache, [0, 1], [2, 3]) # Pull data from blocks 0 and 1 on the Prefill side to blocks 2 and 3 on the Decode side.

On the Prefill side, call the transfer_cache_async API to transmit data to the Decode side.

# Implement LayerSynchronizerImpl to obtain the computation completion status of each layer via torch Event. In this example, the Event mechanism is used.
class LayerSynchronizerImpl(LayerSynchronizer):
    def __init__(self, events):
        self._events = events

    def synchronize_layer(self, layer_index: int, timeout_in_millis: Optional[int]) -> bool:
        self._events[layer_index].wait()
        return True


events = [torch.npu.Event() for _ in range(cache_desc.num_tensors // 2)]
# Execute the model. The model calls events[layer_index].record() after the computation of each layer is complete.
# This function is implemented by users.
user_model.Predict(kv_cache_tensors, events)

# After the model execution is initiated, call transfer_cache_async to transmit data. The memory addresses of the KV cache tensors for each layer allocated by the Decode side need to be provided here.
transfer_config = TransferConfig(DECODER_CLUSTER_ID, decoder_kv_cache_addrs)
cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(events), [transfer_config], [0, 1],
                                                   [2, 3])
# Wait for the transmission result synchronously.
cache_task.synchronize()

After the service is complete, call deallocate_cache on the Prefill and Decode sides to release the allocated KV cache memory.

# Wait for the Decode side to finish pulling the KV cache for the corresponding request.
# Release the KV cache memory for the corresponding request at the appropriate time based on the service logic. In the PagedAttention scenario, there is no need to release the corresponding cache_key.
kv_cache_manager.deallocate_cache(kv_cache)

When the service exits, disconnect the link and call finalize of llm_datadist on the Prefill and Decode sides based on the cluster link disconnection example.

Troubleshooting

LLM_DEVICE_OUT_OF_MEMORY indicates insufficient device memory for KV cache allocation. To address this, check the value of ge.flowGraphMemMaxSize 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 whether 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 occurs when a KV cache pull operation times out, indicating a link issue. To address this, try disconnecting and re-establishing the link.

Parent topic: Functions