allocate_cache

Applicability

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

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 Decode side successfully calls the pull_cache or push_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

from llm_datadist import *
...
cache_desc = CacheDesc(1, [2, 1024 * 1024], DataType.DT_FLOAT16)
cache_keys = [CacheKey(1, req_id=1), CacheKey(1, req_id=2)]
cache = cache_manager.allocate_cache(cache_desc, cache_keys)

Returns

In normal cases, Cache is returned.

If the input data type is incorrect, the TypeError or ValueError exception is reported.

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.

Restrictions

• 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.
• The memory pool must be configured before this API call.