transfer_cache_async

Applicable Products

Product	Supported (Yes/No)
Ascend 950PR/Ascend 950DT	No
Atlas A3 training products/Atlas A3 inference products	Yes
Atlas A2 training products/Atlas A2 inference products	Yes
Atlas 200I/500 A2 inference products	No
Atlas inference products	No
Atlas training products	No

Note: For the Atlas A2 training products/Atlas A2 inference products, only the Atlas 800I A2 inference server and A200I A2 Box heterogeneous subrack are supported.

Function Description

Asynchronously transfers the KV cache in a hierarchical manner.

Prototype

transfer_cache_async(src_cache: KvCache,
                     layer_synchronizer: LayerSynchronizer,
                     transfer_configs: Union[List[TransferConfig], Tuple[TransferConfig]],
                     src_block_indices: Optional[Union[List[int], Tuple[int]]] = None,
                     dst_block_indices: Optional[Union[List[int], Tuple[int]]] = None,
                     dst_block_memory_size: Optional[int] = None) -> CacheTask

Parameters

Parameter	Data Type	Value Description
src_cache	KvCache	Source cache.
layer_synchronizer	LayerSynchronizer	LayerSynchronizer implementation class object.
transfer_configs	Union[List[TransferConfig], Tuple[TransferConfig]]	Transmission configuration list or tuple.
src_block_indices	Optional[Union[List[int], Tuple[int]]]	Block indices of the source cache, which are set when the source cache is in the PagedAttention scenario.
dst_block_indices	Optional[Union[List[int], Tuple[int]]]	Block indices of the destination cache, which are set when the destination cache is in the PagedAttention scenario.
dst_block_memory_size	Optional[int]	Size of the memory occupied by each block in the destination cache, which is set when the destination cache is in the PagedAttention scenario. If the source cache is also in the PagedAttention scenario, this parameter can be omitted. In this case, it is automatically set to the memory size occupied by each block in the source cache. This parameter is omitted if the value is set to 0.

Example

from llm_datadist import *
...
class LayerSynchronizerImpl(LayerSynchronizer):
    def synchronize_layer(self, layer_index: int, timeout_in_millis: Optional[int]) -> bool:
        # need control time for transfer layer here.
        return True
num_layers = 40
dst_cluster_id = 2
# need register decoder kv addr here.
decoder_addrs = ...
assert(len(decoder_addrs) = 2*num_layers)
transfer_config = TransferConfig(dst_cluster_id, decoder_addrs, range(0, num_layers), 0)
cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(), [transfer_config])
cache_task.synchronize()
cache_task.get_results()

Returns

In normal cases, a CacheTask instance is returned.

If the input data type is incorrect, a TypeError or ValueError is thrown.

If the input data is invalid, an LLMException is thrown.

Constraints

Currently, only the scenarios where both src_cache and dst_cache are continuous caches or both are PA caches are supported.
This API cannot be used concurrently with pull_cache and pull_blocks over the same link.
This API does not support concurrent calls.
This API cannot be called in single-process multi-device mode.

Parent topic: KvCacheManager