Documentation
Download

KV Cache Management

Function Description

KV cache management involves the following APIs and functions.

Table 1 Main KV cache management APIs and their functions

API

Function

register_cache

Registers a cache.

In the non-PagedAttention scenario, this API is called to register a self-allocated memory.

register_blocks_cache

Registers a cache.

In the PagedAttention scenario, this API is called to register a self-allocated memory.

unregister_cache

Deregisters a registered cache when it is no longer used.

pull_cache

Pulls a KV cache from the remote node to the local KV cache based on CacheKey.

CacheKey must be the same as that in allocate_cache.

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.

transfer_cache_async

Asynchronously transmits the KV cache in a hierarchical manner.

push_blocks

Pushes a KV cache from the local node to the remote KV cache in the PagedAttention scenario.

push_cache

Pushes a KV cache from the local node to the remote KV cache in the non-PagedAttention scenario.

Scenario

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

  1. Initialize the LLM-DataDist instance on the Prefill and Decode sides as described in Link Management.
  2. Allocate the KV cache at each layer of the model on the Prefill and Decode sides in advance based on the calculated size. Different requests reuse the allocated KV cache. The upper-layer framework manages this process, and the allocated memory is released when the service ends.
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
    import torchair
import torch
import torch_npu
# Obtain cache_manager from the initialized llm_datadist.
cache_manager = llm_datadist.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)
tensor1 = torch.full((4, 4, 8), 1, dtype=torch.float).npu()
... # Allocate other tensors.
cache = cache_manager.register_cache(cache_desc, [int(tensor.data_ptr()), int(tensor2.data_ptr()) ...])

 # After link establishment, pass the registered kv_tensors to the model for inference computation to generate the KV cache, and transmit the model output to the incremental inference model as input.
  3. Establish a link between the LLM-DataDist instances on Prefill and Decode sides as described in Link Management.
  4. 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.
      1
2
3
      # 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)
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.
       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
      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 = cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(events), [transfer_config])
# Wait for the transmission result synchronously.
cache_task.synchronize()
  5. Taking PyTorch as an example, convert the KV cache into torch tensors and perform incremental model inference.
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
    # 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 for the request's incremental inference to complete.
  6. Deregister the KV cache memory for the corresponding request at the appropriate time based on the service logic.
    1
    cache_manager.unregister_cache(cache_id)
  7. 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 registration, deregistration, and transmission. Pseudocode examples are provided below for each role.

  1. Initialize the LLM-DataDist instance on the Prefill and Decode sides as described in Link Management.
  2. Allocate tensors, such as torch tensors, at each layer of the model on the Prefill and Decode sides based on the calculated num_blocks, and register the tensors with the LLM-DataDist instance. 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.
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
    # Prefill/Decode side
# Obtain kv_cache_manager from the initialized llm_datadist.
cache_manager = llm_datadist.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 the cluster_id used during llm_datadist initialization.
cache_key = BlocksCacheKey(prompt_cluster_id=0, model_id=0)
... # Allocate tensors.
# Call register_blocks_cache to register the KV cache.
kv_cache = cache_manager.register_blocks_cache(cache_desc, [addr, addr2, ...(tensor address)], cache_key)
  3. Establish a link between the LLM-DataDist instances on Prefill and Decode sides as described in Link Management.
  4. 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.
  5. 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.
      1
2
3
4
      # 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 the 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 register_blocks_cache is called on the Prefill side.
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.
       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
      # 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 = cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(events), [transfer_config], [0, 1], [2, 3])
# Wait for the transmission result synchronously.
cache_task.synchronize()
  6. After the service ends, call unregister_cache on the Prefill and Decode sides to deregister the KV cache memory.
    1
2
3
    # Wait for the Decode side to finish pulling the KV cache for the corresponding request.
# Deregister the KV cache for the corresponding request at the appropriate time based on the service logic.
cache_manager.unregister_cache(cache_id)
  7. When the service exits, disconnect the links on the Prefill and Decode sides by following the example of cluster disconnection and call finalize to release 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 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: Function Description
This site uses cookies to improve your browsing experience. By continuing to browse this site, you agree to our use of cookies.View Details