allocate_cache

Applicability

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

Allocates a cache. After a cache is successfully allocated, it is referenced by cache_id and cache_keys at the same time. The resources occupied by the cache are released only after these references are removed.

The reference of cache_id needs to be removed by using deallocate_cache. The reference of cache_keys can be removed by using either of the following methods:

• The Decoder side successfully calls the pull_cache API.
• The Prompt side calls the remove_cache_key API.

Prototype

1

allocate_cache(cache_desc: CacheDesc, cache_keys: Union[Tuple[CacheKey], List[CacheKey]] = ())

Parameters

Example

1
2
3
4
5
6

from llm_datadist import *
...
kv_cache_manager = data_dist.kv_cache_manager
cache_desc = CacheDesc(80, [2, 2 * 1024 * 1024], DataType.DT_FLOAT16)
cache_keys = [CacheKey(prompt_cluster_id=0, req_id=1)]
kv_cache = kv_cache_manager.allocate_cache(cache_desc, cache_keys)

Returns

In normal cases, a KvCache instance is returned.

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

If cache_keys contains the cache key bound during memory allocation, an LLMException is thrown.

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

Constraints

• When cache_keys is passed in and the batch size is greater than 1, you must provide the same number of cache keys, each referencing a group of KV tensors.
• If the current inference batch is not fully occupied (that is, invalid batch indices exist), insert a special CacheKey (with req_id set to UINT64_MAX) as a placeholder for each invalid index. Placeholders for idle batch indices at the end of the batch can be omitted.
• If duplicate cache keys exist, the last one takes effect.
• This API does not support concurrent calls.