allocate_blocks_cache

Applicable Products

Product	Supported (Yes/No)
Ascend 950PR/Ascend 950DT	No
Atlas A3 training products/Atlas A3 inference products	Yes
Atlas A2 training products/Atlas A2 inference products	Yes
Atlas 200I/500 A2 inference products	No
Atlas inference products	No
Atlas training products	No

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

Function Description

In the PagedAttention scenario, allocates the cache with multiple blocks. After the cache is successfully allocated, the memory can be released through deallocate_cache.

Prototype

allocate_blocks_cache(cache_desc: CacheDesc, blocks_cache_key: Optional[BlocksCacheKey] = None)

Parameters

Parameter	Data Type	Value Description
cache_desc	CacheDesc	Cache description.
blocks_cache_key	Optional[BlocksCacheKey]	This parameter can be set only when LLMRole is set to PROMPT, and it is used for the Decode side to pull KV cache.

Example

from llm_datadist import BlocksCacheKey
num_blocks = 1000
cache_desc = CacheDesc(80, [num_blocks , 128 * 1024], DataType.DT_FLOAT16)
blocks_cache_key = BlocksCacheKey(0, 0)
kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc, blocks_cache_key)

Returns

In normal cases, a KvCache instance 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 does not support concurrent calls.

Parent topic: KvCacheManager