pull_blocks

Applicability

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

Note: For 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

In the PagedAttention scenario, the cache is pulled from the peer node in block list based on BlocksCacheKey to the local cache.

Prototype

pull_blocks(src_cache_key: Union[CacheKey, CacheKeyByIdAndIndex, BlocksCacheKey],
            dst_cache: Cache,
            src_blocks: Optional[Union[Tuple[int], List[int]]] = (),
            dst_blocks: Union[Tuple[int], List[int]] = (), **kwargs)

Parameters

Parameter	Data Type	Description
src_cache_key	Union[CacheKey, CacheKeyByIdAndIndex, BlocksCacheKey]	Remote cache index.
dst_cache	Cache	Target cache.
src_blocks	Optional[Union[Tuple[int], List[int]]]	Remote block index list. If src_cache_key is not BlocksCacheKey, leave this field blank.
dst_blocks	Union[Tuple[int], List[int]]	Local block index list.
**kwargs	N/A	This is the typical way to handle extensible parameters in Python functions: parameters are passed in using the key=value format. For details about optional parameters, see Table 1.

**Table 1** Optional parameters of **kwargs
Parameter	Data Type	Description
src_layer_range	Optional[range]	(Optional) Used for pulling KVs by layer. Layer range of the transmission source. The step can only be 1. If this parameter is not set, data at all layers is transmitted. Note that this is the index of the layer, not the index of the tensor. That is, one layer corresponds to N contiguous tensors (K/V). If memory allocation is required, the tensors must be arranged in KV order. Other arrangements are not supported. N is the value of tensor_num_per_layer. The default value is 2.
dst_layer_range	Optional[range]	(Optional) Used for pulling KVs by layer. Layer range of the transmission destination. The step can only be 1. If this parameter is not set, data at all layers is transmitted. Note that this is the index of the layer, not the index of the tensor. That is, one layer corresponds to N contiguous tensors (K/V). If memory allocation is required, the tensors must be arranged in KV order. Other arrangements are not supported. N is the value of tensor_num_per_layer. The default value is 2.
tensor_num_per_layer	Optional[int]	(Optional) Number of tensors at each layer. The default value is 2. The value range is [1, total number of tensors in the cache]. When src_layer_range or dst_layer_range uses non-default values, tensor_num_per_layer can either remain at its default value or be set to another value, which must be divisible by the total number of tensors in the cache.

Example

Click GitCode, select the matching version, and obtain the sample from the examples/python directory.

Returns

In normal cases, no value is returned.
If the input data type is incorrect, the TypeError or ValueError exception is reported.
If the execution time exceeds the value of sync_kv_timeout, an LLMException is thrown.
If the layer_range parameter is abnormal, an LLMException is thrown.

Restrictions

When src_cache is on the host and dst_cache is on the device, only scenarios where both caches are block caches are supported.
Enabling enable_remote_cache_accessible introduces additional restrictions. For details, see the restrictions description for this function.
In D2H and H2D transmission scenarios, the device memory pool must be configured during host initialization, with the memory pool size set to at least 100 MB.

Parent topic: CacheManager