pull_blocks

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

Pulls the KV cache from the corresponding Prompt node to the local KV cache via a block list based on BlocksCacheKey in the PagedAttention scenario. This API can be called only when LLMRole is set to DECODER.

Prototype

pull_blocks(prompt_cache_key: BlocksCacheKey, decoder_kv_cache: KvCache, prompt_blocks: List[int], decoder_blocks: List[int], **kwargs)

Parameters

Parameter	Data Type	Description
prompt_cache_key	BlocksCacheKey	BlocksCacheKey to be pulled.
decoder_kv_cache	KvCache	Target KV cache.
prompt_blocks	List[int]	Block index list of the Prompt node.
decoder_blocks	List[int]	Block index list of the Decode node.
**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

from llm_datadist import *
...
kv_cache_manager.pull_blocks(prompt_cache_key, kv_cache, [0, 1], [2, 3])

# Example of enabling the layer_range function
kv_cache_manager.pull_blocks(prompt_cache_key, kv_cache, [0, 1], [2, 3], src_layer_range=range(2), dst_layer_range=range(2))

Returns

In normal cases, no value is returned.

If a parameter is incorrect, a TypeError or ValueError may be thrown.

If the execution time exceeds the value of sync_kv_timeout, an LLMException is thrown.

Constraints

This API cannot be used concurrently with transfer_cache_async over the same link.
This API does not support concurrent calls by multiple threads.

Parent topic: KvCacheManager