pull_cache
Applicability
Product |
Supported (Yes/No) |
|---|---|
Atlas 350 Accelerator Card |
No |
Yes |
|
Yes |
|
No |
|
No |
|
No |
Note: For the
Function Description
Pulls the KV cache from the corresponding Prompt node to the local KV cache based on CacheKey. This API can be called only when LLMRole is set to DECODER.
Prototype
1 | pull_cache(cache_key: Union[CacheKey, CacheKeyByIdAndIndex], kv_cache: KvCache, batch_index: int = 0, size: int = -1, **kwargs) |
Parameters
Parameter |
Data Type |
Value Description |
|---|---|---|
cache_key |
Union[CacheKey, CacheKeyByIdAndIndex] |
CacheKey to be pulled. The value of CacheKey must be the same as that of CacheKey in allocate_cache. Pass CacheKey to obtain the value by req_id, prefix_id, or model_id. Pass CacheKeyByIdAndIndex to obtain the value by cache_id or batch_index. |
kv_cache |
Target KV cache. |
|
batch_index |
int |
Batch index of the target KV cache. Defaults to 0. |
size |
int |
Defaults to -1. Set this parameter to an integer greater than 0, indicating 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 calculated as the size of a single local KV entry minus the value of dst_cache_offset. |
**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. |
Parameter |
Data Type |
Value 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. |
src_cache_offset |
Optional[int] |
Set it to an integer greater or equal to 0. Offset of the transmission source from which data of the specified size is to be pulled. |
dst_cache_offset |
Optional[int] |
Set it to an integer greater or equal to 0. Offset of the transmission destination to which data is to be pull. |
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
1 2 3 4 5 6 7 8 | from llm_datadist import * ... cache_keys = [CacheKey(1, req_id=1), CacheKey(1, req_id=2)] kv_cache_manager.pull_cache(cache_keys[0], cache, 0) # Enable the layer_range function. kv_cache_manager.pull_cache(cache_keys[1], cache, 1, src_layer_range=range(0,2), dst_layer_range=range(2,4)) # Enable the offset function. kv_cache_manager.pull_cache(cache_keys[1], cache, src_cache_offset=0, dst_cache_offset=0) |
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.