transfer_cache_async

Applicability

Product	Supported (Yes/No)
Atlas 350 Accelerator Card	No
Atlas A3 training product/Atlas A3 inference product	Yes
Atlas A2 training product/Atlas A2 inference product	Yes
Atlas 200I/500 A2 inference product	No
Atlas inference product	No
Atlas training product	No

Note: For the Atlas A2 training product/Atlas A2 inference product, 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	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