pull_cache

Applicability

Product	Supported
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	x
Atlas inference products	x
Atlas training products	x

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

Function Description

Pulls the cache from the peer node to the local cache based on the cache key.

Prototype

pull_cache(cache_key: Union[CacheKey, CacheKeyByIdAndIndex],
           cache: Cache,
           batch_index: int = 0,
           size: int = -1, **kwargs)

Parameters

Parameter	Data Type	Description
cache_key	Union[CacheKey, CacheKeyByIdAndIndex]	Key of the cache to be pulled. Pass CacheKey to pull the cache via req_id, prefix_id, and model_id. Pass CacheKeyByIdAndIndex to pull the cache via cache_id and batch_index.
cache	Cache	Target cache.
batch_index	int	Batch index. The default value is 0.
size	int	If this parameter is set to an integer greater than 0, it indicates the size of the tensor to be pulled. If this parameter is set to -1, it indicates a full copy, where the actual size is the size of a single local KV entry. The default value is -1.
**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) Layer range of the transmission source, which is used for pulling KVs by layer. 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 consecutive 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, which defaults to 2.
dst_layer_range	Optional[range]	(Optional) Layer range of the transmission destination, which is used for pulling KVs by layer. 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 consecutive 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, which defaults to 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, a TypeError or ValueError is thrown.

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 enable_remote_cache_accessible is set to True, only the cache_key type CacheKeyByIdAndIndex is supported.
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